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