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.ImageReadException;
029import org.apache.commons.imaging.ImageWriteException;
030import org.apache.commons.imaging.common.BinaryFileParser;
031import org.apache.commons.imaging.common.ByteConversions;
032import org.apache.commons.imaging.common.bytesource.ByteSource;
033import org.apache.commons.imaging.formats.jpeg.JpegConstants;
034import org.apache.commons.imaging.formats.jpeg.JpegUtils;
035import org.apache.commons.imaging.formats.jpeg.iptc.IptcParser;
036
037/**
038 * Interface for Exif write/update/remove functionality for Jpeg/JFIF images.
039 */
040public class JpegRewriter extends BinaryFileParser {
041    private static final ByteOrder JPEG_BYTE_ORDER = ByteOrder.BIG_ENDIAN;
042    private static final SegmentFilter EXIF_SEGMENT_FILTER = segment -> segment.isExifSegment();
043    private static final SegmentFilter XMP_SEGMENT_FILTER = segment -> segment.isXmpSegment();
044    private static final SegmentFilter PHOTOSHOP_APP13_SEGMENT_FILTER = segment -> segment.isPhotoshopApp13Segment();
045
046    /**
047     * Constructor. to guess whether a file contains an image based on its file
048     * extension.
049     */
050    public JpegRewriter() {
051        setByteOrder(JPEG_BYTE_ORDER);
052    }
053
054    protected static class JFIFPieces {
055        public final List<JFIFPiece> pieces;
056        public final List<JFIFPiece> segmentPieces;
057
058        public JFIFPieces(final List<JFIFPiece> pieces,
059                final List<JFIFPiece> segmentPieces) {
060            this.pieces = pieces;
061            this.segmentPieces = segmentPieces;
062        }
063
064    }
065
066    protected abstract static class JFIFPiece {
067        protected abstract void write(OutputStream os) throws IOException;
068
069        @Override
070        public String toString() {
071            return "[" + this.getClass().getName() + "]";
072        }
073    }
074
075    protected static class JFIFPieceSegment extends JFIFPiece {
076        public final int marker;
077        private final byte[] markerBytes;
078        private final byte[] segmentLengthBytes;
079        private final byte[] segmentData;
080
081        public JFIFPieceSegment(final int marker, final byte[] segmentData) {
082            this(marker,
083                    ByteConversions.toBytes((short) marker, JPEG_BYTE_ORDER),
084                    ByteConversions.toBytes((short) (segmentData.length + 2), JPEG_BYTE_ORDER),
085                    segmentData);
086        }
087
088        JFIFPieceSegment(final int marker, final byte[] markerBytes,
089                final byte[] segmentLengthBytes, final byte[] segmentData) {
090            this.marker = marker;
091            this.markerBytes = markerBytes;
092            this.segmentLengthBytes = segmentLengthBytes;
093            this.segmentData = segmentData.clone();
094        }
095
096        @Override
097        public String toString() {
098            return "[" + this.getClass().getName() + " (0x"
099                    + Integer.toHexString(marker) + ")]";
100        }
101
102        @Override
103        protected void write(final OutputStream os) throws IOException {
104            os.write(markerBytes);
105            os.write(segmentLengthBytes);
106            os.write(segmentData);
107        }
108
109        public boolean isApp1Segment() {
110            return marker == JpegConstants.JPEG_APP1_MARKER;
111        }
112
113        public boolean isAppSegment() {
114            return marker >= JpegConstants.JPEG_APP0_MARKER && marker <= JpegConstants.JPEG_APP15_MARKER;
115        }
116
117        public boolean isExifSegment() {
118            if (marker != JpegConstants.JPEG_APP1_MARKER) {
119                return false;
120            }
121            if (!startsWith(segmentData, JpegConstants.EXIF_IDENTIFIER_CODE)) {
122                return false;
123            }
124            return true;
125        }
126
127        public boolean isPhotoshopApp13Segment() {
128            if (marker != JpegConstants.JPEG_APP13_MARKER) {
129                return false;
130            }
131            if (!new IptcParser().isPhotoshopJpegSegment(segmentData)) {
132                return false;
133            }
134            return true;
135        }
136
137        public boolean isXmpSegment() {
138            if (marker != JpegConstants.JPEG_APP1_MARKER) {
139                return false;
140            }
141            if (!startsWith(segmentData, JpegConstants.XMP_IDENTIFIER)) {
142                return false;
143            }
144            return true;
145        }
146
147        public byte[] getSegmentData() {
148            return segmentData.clone();
149        }
150
151    }
152
153    static class JFIFPieceImageData extends JFIFPiece {
154        private final byte[] markerBytes;
155        private final byte[] imageData;
156
157        JFIFPieceImageData(final byte[] markerBytes, final byte[] imageData) {
158            super();
159            this.markerBytes = markerBytes;
160            this.imageData = imageData;
161        }
162
163        @Override
164        protected void write(final OutputStream os) throws IOException {
165            os.write(markerBytes);
166            os.write(imageData);
167        }
168    }
169
170    protected JFIFPieces analyzeJFIF(final ByteSource byteSource) throws ImageReadException, IOException {
171        final List<JFIFPiece> pieces = new ArrayList<>();
172        final List<JFIFPiece> segmentPieces = new ArrayList<>();
173
174        final JpegUtils.Visitor visitor = new JpegUtils.Visitor() {
175            // return false to exit before reading image data.
176            @Override
177            public boolean beginSOS() {
178                return true;
179            }
180
181            @Override
182            public void visitSOS(final int marker, final byte[] markerBytes, final byte[] imageData) {
183                pieces.add(new JFIFPieceImageData(markerBytes, imageData));
184            }
185
186            // return false to exit traversal.
187            @Override
188            public boolean visitSegment(final int marker, final byte[] markerBytes,
189                    final int segmentLength, final byte[] segmentLengthBytes,
190                    final byte[] segmentData) throws ImageReadException, IOException {
191                final JFIFPiece piece = new JFIFPieceSegment(marker, markerBytes,
192                        segmentLengthBytes, segmentData);
193                pieces.add(piece);
194                segmentPieces.add(piece);
195
196                return true;
197            }
198        };
199
200        new JpegUtils().traverseJFIF(byteSource, visitor);
201
202        return new JFIFPieces(pieces, segmentPieces);
203    }
204
205    private interface SegmentFilter {
206        boolean filter(JFIFPieceSegment segment);
207    }
208
209    protected <T extends JFIFPiece> List<T> removeXmpSegments(final List<T> segments) {
210        return filterSegments(segments, XMP_SEGMENT_FILTER);
211    }
212
213    protected <T extends JFIFPiece> List<T> removePhotoshopApp13Segments(
214            final List<T> segments) {
215        return filterSegments(segments, PHOTOSHOP_APP13_SEGMENT_FILTER);
216    }
217
218    protected <T extends JFIFPiece> List<T> findPhotoshopApp13Segments(
219            final List<T> segments) {
220        return filterSegments(segments, PHOTOSHOP_APP13_SEGMENT_FILTER, true);
221    }
222
223    protected <T extends JFIFPiece> List<T> removeExifSegments(final List<T> segments) {
224        return filterSegments(segments, EXIF_SEGMENT_FILTER);
225    }
226
227    protected <T extends JFIFPiece> List<T> filterSegments(final List<T> segments,
228            final SegmentFilter filter) {
229        return filterSegments(segments, filter, false);
230    }
231
232    protected <T extends JFIFPiece> List<T> filterSegments(final List<T> segments,
233            final SegmentFilter filter, final boolean reverse) {
234        final List<T> result = new ArrayList<>();
235
236        for (final T piece : segments) {
237            if (piece instanceof JFIFPieceSegment) {
238                if (filter.filter((JFIFPieceSegment) piece) ^ !reverse) {
239                    result.add(piece);
240                }
241            } else if (!reverse) {
242                result.add(piece);
243            }
244        }
245
246        return result;
247    }
248
249    protected <T extends JFIFPiece, U extends JFIFPiece> List<JFIFPiece> insertBeforeFirstAppSegments(
250            final List<T> segments, final List<U> newSegments) throws ImageWriteException {
251        int firstAppIndex = -1;
252        for (int i = 0; i < segments.size(); i++) {
253            final JFIFPiece piece = segments.get(i);
254            if (!(piece instanceof JFIFPieceSegment)) {
255                continue;
256            }
257
258            final JFIFPieceSegment segment = (JFIFPieceSegment) piece;
259            if (segment.isAppSegment()) {
260                if (firstAppIndex == -1) {
261                    firstAppIndex = i;
262                }
263            }
264        }
265
266        final List<JFIFPiece> result = new ArrayList<>(segments);
267        if (firstAppIndex == -1) {
268            throw new ImageWriteException("JPEG file has no APP segments.");
269        }
270        result.addAll(firstAppIndex, newSegments);
271        return result;
272    }
273
274    protected <T extends JFIFPiece, U extends JFIFPiece> List<JFIFPiece> insertAfterLastAppSegments(
275            final List<T> segments, final List<U> newSegments) throws ImageWriteException {
276        int lastAppIndex = -1;
277        for (int i = 0; i < segments.size(); i++) {
278            final JFIFPiece piece = segments.get(i);
279            if (!(piece instanceof JFIFPieceSegment)) {
280                continue;
281            }
282
283            final JFIFPieceSegment segment = (JFIFPieceSegment) piece;
284            if (segment.isAppSegment()) {
285                lastAppIndex = i;
286            }
287        }
288
289        final List<JFIFPiece> result = new ArrayList<>(segments);
290        if (lastAppIndex == -1) {
291            if (segments.isEmpty()) {
292                throw new ImageWriteException("JPEG file has no APP segments.");
293            }
294            result.addAll(1, newSegments);
295        } else {
296            result.addAll(lastAppIndex + 1, newSegments);
297        }
298
299        return result;
300    }
301
302    protected void writeSegments(final OutputStream outputStream,
303            final List<? extends JFIFPiece> segments) throws IOException {
304        try (DataOutputStream os = new DataOutputStream(outputStream)) {
305            JpegConstants.SOI.writeTo(os);
306
307            for (final JFIFPiece piece : segments) {
308                piece.write(os);
309            }
310        }
311    }
312
313    // private void writeSegment(OutputStream os, JFIFPieceSegment piece)
314    // throws ImageWriteException, IOException
315    // {
316    // byte markerBytes[] = convertShortToByteArray(JPEG_APP1_MARKER,
317    // JPEG_BYTE_ORDER);
318    // if (piece.segmentData.length > 0xffff)
319    // throw new JpegSegmentOverflowException("Jpeg segment is too long: "
320    // + piece.segmentData.length);
321    // int segmentLength = piece.segmentData.length + 2;
322    // byte segmentLengthBytes[] = convertShortToByteArray(segmentLength,
323    // JPEG_BYTE_ORDER);
324    //
325    // os.write(markerBytes);
326    // os.write(segmentLengthBytes);
327    // os.write(piece.segmentData);
328    // }
329
330    public static class JpegSegmentOverflowException extends ImageWriteException {
331        private static final long serialVersionUID = -1062145751550646846L;
332
333        public JpegSegmentOverflowException(final String message) {
334            super(message);
335        }
336    }
337
338}