001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.commons.imaging.formats.jpeg.xmp;
018
019import static org.apache.commons.imaging.common.BinaryFunctions.startsWith;
020
021import java.io.DataOutputStream;
022import java.io.IOException;
023import java.io.OutputStream;
024import java.nio.ByteOrder;
025import java.util.ArrayList;
026import java.util.List;
027
028import org.apache.commons.imaging.ImagingException;
029import org.apache.commons.imaging.bytesource.ByteSource;
030import org.apache.commons.imaging.common.BinaryFileParser;
031import org.apache.commons.imaging.common.ByteConversions;
032import org.apache.commons.imaging.formats.jpeg.JpegConstants;
033import org.apache.commons.imaging.formats.jpeg.JpegUtils;
034import org.apache.commons.imaging.formats.jpeg.iptc.IptcParser;
035
036/**
037 * Interface for Exif write/update/remove functionality for Jpeg/JFIF images.
038 */
039public class JpegRewriter extends BinaryFileParser {
040    protected abstract static class JFIFPiece {
041        @Override
042        public String toString() {
043            return "[" + this.getClass().getName() + "]";
044        }
045
046        protected abstract void write(OutputStream os) throws IOException;
047    }
048
049    static class JFIFPieceImageData extends JFIFPiece {
050        private final byte[] markerBytes;
051        private final byte[] imageData;
052
053        JFIFPieceImageData(final byte[] markerBytes, final byte[] imageData) {
054            this.markerBytes = markerBytes;
055            this.imageData = imageData;
056        }
057
058        @Override
059        protected void write(final OutputStream os) throws IOException {
060            os.write(markerBytes);
061            os.write(imageData);
062        }
063    }
064
065    protected static class JFIFPieces {
066        public final List<JFIFPiece> pieces;
067        public final List<JFIFPiece> segmentPieces;
068
069        public JFIFPieces(final List<JFIFPiece> pieces, final List<JFIFPiece> segmentPieces) {
070            this.pieces = pieces;
071            this.segmentPieces = segmentPieces;
072        }
073
074    }
075
076    protected static class JFIFPieceSegment extends JFIFPiece {
077        public final int marker;
078        private final byte[] markerBytes;
079        private final byte[] segmentLengthBytes;
080        private final byte[] segmentData;
081
082        public JFIFPieceSegment(final int marker, final byte[] segmentData) {
083            this(marker, ByteConversions.toBytes((short) marker, JPEG_BYTE_ORDER), ByteConversions.toBytes((short) (segmentData.length + 2), JPEG_BYTE_ORDER),
084                    segmentData);
085        }
086
087        JFIFPieceSegment(final int marker, final byte[] markerBytes, final byte[] segmentLengthBytes, final byte[] segmentData) {
088            this.marker = marker;
089            this.markerBytes = markerBytes;
090            this.segmentLengthBytes = segmentLengthBytes;
091            this.segmentData = segmentData.clone();
092        }
093
094        public byte[] getSegmentData() {
095            return segmentData.clone();
096        }
097
098        public boolean isApp1Segment() {
099            return marker == JpegConstants.JPEG_APP1_MARKER;
100        }
101
102        public boolean isAppSegment() {
103            return marker >= JpegConstants.JPEG_APP0_MARKER && marker <= JpegConstants.JPEG_APP15_MARKER;
104        }
105
106        public boolean isExifSegment() {
107            if (marker != JpegConstants.JPEG_APP1_MARKER) {
108                return false;
109            }
110            if (!startsWith(segmentData, JpegConstants.EXIF_IDENTIFIER_CODE)) {
111                return false;
112            }
113            return true;
114        }
115
116        public boolean isPhotoshopApp13Segment() {
117            if (marker != JpegConstants.JPEG_APP13_MARKER) {
118                return false;
119            }
120            if (!new IptcParser().isPhotoshopJpegSegment(segmentData)) {
121                return false;
122            }
123            return true;
124        }
125
126        public boolean isXmpSegment() {
127            if (marker != JpegConstants.JPEG_APP1_MARKER) {
128                return false;
129            }
130            if (!startsWith(segmentData, JpegConstants.XMP_IDENTIFIER)) {
131                return false;
132            }
133            return true;
134        }
135
136        @Override
137        public String toString() {
138            return "[" + this.getClass().getName() + " (0x" + Integer.toHexString(marker) + ")]";
139        }
140
141        @Override
142        protected void write(final OutputStream os) throws IOException {
143            os.write(markerBytes);
144            os.write(segmentLengthBytes);
145            os.write(segmentData);
146        }
147
148    }
149
150    private interface SegmentFilter {
151        boolean filter(JFIFPieceSegment segment);
152    }
153
154    private static final ByteOrder JPEG_BYTE_ORDER = ByteOrder.BIG_ENDIAN;
155
156    private static final SegmentFilter EXIF_SEGMENT_FILTER = JFIFPieceSegment::isExifSegment;
157
158    private static final SegmentFilter XMP_SEGMENT_FILTER = JFIFPieceSegment::isXmpSegment;
159
160    private static final SegmentFilter PHOTOSHOP_APP13_SEGMENT_FILTER = JFIFPieceSegment::isPhotoshopApp13Segment;
161
162    /**
163     * Constructs a new instance. to guess whether a file contains an image based on its file extension.
164     */
165    public JpegRewriter() {
166        super(JPEG_BYTE_ORDER);
167    }
168
169    protected JFIFPieces analyzeJfif(final ByteSource byteSource) throws ImagingException, IOException {
170        final List<JFIFPiece> pieces = new ArrayList<>();
171        final List<JFIFPiece> segmentPieces = new ArrayList<>();
172
173        final JpegUtils.Visitor visitor = new JpegUtils.Visitor() {
174            // return false to exit before reading image data.
175            @Override
176            public boolean beginSos() {
177                return true;
178            }
179
180            // return false to exit traversal.
181            @Override
182            public boolean visitSegment(final int marker, final byte[] markerBytes, final int segmentLength, final byte[] segmentLengthBytes,
183                    final byte[] segmentData) throws ImagingException, IOException {
184                final JFIFPiece piece = new JFIFPieceSegment(marker, markerBytes, segmentLengthBytes, segmentData);
185                pieces.add(piece);
186                segmentPieces.add(piece);
187
188                return true;
189            }
190
191            @Override
192            public void visitSos(final int marker, final byte[] markerBytes, final byte[] imageData) {
193                pieces.add(new JFIFPieceImageData(markerBytes, imageData));
194            }
195        };
196
197        new JpegUtils().traverseJfif(byteSource, visitor);
198
199        return new JFIFPieces(pieces, segmentPieces);
200    }
201
202    protected <T extends JFIFPiece> List<T> filterSegments(final List<T> segments, final SegmentFilter filter) {
203        return filterSegments(segments, filter, false);
204    }
205
206    protected <T extends JFIFPiece> List<T> filterSegments(final List<T> segments, final SegmentFilter filter, final boolean reverse) {
207        final List<T> result = new ArrayList<>();
208
209        for (final T piece : segments) {
210            if (piece instanceof JFIFPieceSegment) {
211                if (filter.filter((JFIFPieceSegment) piece) == reverse) {
212                    result.add(piece);
213                }
214            } else if (!reverse) {
215                result.add(piece);
216            }
217        }
218
219        return result;
220    }
221
222    protected <T extends JFIFPiece> List<T> findPhotoshopApp13Segments(final List<T> segments) {
223        return filterSegments(segments, PHOTOSHOP_APP13_SEGMENT_FILTER, true);
224    }
225
226    protected <T extends JFIFPiece, U extends JFIFPiece> List<JFIFPiece> insertAfterLastAppSegments(final List<T> segments, final List<U> newSegments)
227            throws ImagingException {
228        int lastAppIndex = -1;
229        for (int i = 0; i < segments.size(); i++) {
230            final JFIFPiece piece = segments.get(i);
231            if (!(piece instanceof JFIFPieceSegment)) {
232                continue;
233            }
234
235            final JFIFPieceSegment segment = (JFIFPieceSegment) piece;
236            if (segment.isAppSegment()) {
237                lastAppIndex = i;
238            }
239        }
240
241        final List<JFIFPiece> result = new ArrayList<>(segments);
242        if (lastAppIndex == -1) {
243            if (segments.isEmpty()) {
244                throw new ImagingException("JPEG file has no APP segments.");
245            }
246            result.addAll(1, newSegments);
247        } else {
248            result.addAll(lastAppIndex + 1, newSegments);
249        }
250
251        return result;
252    }
253
254    protected <T extends JFIFPiece, U extends JFIFPiece> List<JFIFPiece> insertBeforeFirstAppSegments(final List<T> segments, final List<U> newSegments)
255            throws ImagingException {
256        int firstAppIndex = -1;
257        for (int i = 0; i < segments.size(); i++) {
258            final JFIFPiece piece = segments.get(i);
259            if (!(piece instanceof JFIFPieceSegment)) {
260                continue;
261            }
262
263            final JFIFPieceSegment segment = (JFIFPieceSegment) piece;
264            if (segment.isAppSegment()) {
265                if (firstAppIndex == -1) {
266                    firstAppIndex = i;
267                }
268            }
269        }
270
271        final List<JFIFPiece> result = new ArrayList<>(segments);
272        if (firstAppIndex == -1) {
273            throw new ImagingException("JPEG file has no APP segments.");
274        }
275        result.addAll(firstAppIndex, newSegments);
276        return result;
277    }
278
279    protected <T extends JFIFPiece> List<T> removeExifSegments(final List<T> segments) {
280        return filterSegments(segments, EXIF_SEGMENT_FILTER);
281    }
282
283    protected <T extends JFIFPiece> List<T> removePhotoshopApp13Segments(final List<T> segments) {
284        return filterSegments(segments, PHOTOSHOP_APP13_SEGMENT_FILTER);
285    }
286
287    protected <T extends JFIFPiece> List<T> removeXmpSegments(final List<T> segments) {
288        return filterSegments(segments, XMP_SEGMENT_FILTER);
289    }
290
291    // private void writeSegment(OutputStream os, JFIFPieceSegment piece)
292    // throws ImageWriteException, IOException
293    // {
294    // byte[] markerBytes = convertShortToByteArray(JPEG_APP1_MARKER,
295    // JPEG_BYTE_ORDER);
296    // if (piece.segmentData.length > 0xffff)
297    // throw new JpegSegmentOverflowException("JPEG segment is too long: "
298    // + piece.segmentData.length);
299    // int segmentLength = piece.segmentData.length + 2;
300    // byte[] segmentLengthBytes = convertShortToByteArray(segmentLength,
301    // JPEG_BYTE_ORDER);
302    //
303    // os.write(markerBytes);
304    // os.write(segmentLengthBytes);
305    // os.write(piece.segmentData);
306    // }
307
308    protected void writeSegments(final OutputStream outputStream, final List<? extends JFIFPiece> segments) throws IOException {
309        try (DataOutputStream os = new DataOutputStream(outputStream)) {
310            JpegConstants.SOI.writeTo(os);
311
312            for (final JFIFPiece piece : segments) {
313                piece.write(os);
314            }
315        }
316    }
317
318}