1package com.googlecode.mp4parser.boxes;
2
3import com.coremedia.iso.Hex;
4import com.coremedia.iso.IsoTypeReader;
5import com.coremedia.iso.IsoTypeWriter;
6import com.coremedia.iso.boxes.Box;
7import com.coremedia.iso.boxes.TrackHeaderBox;
8import com.coremedia.iso.boxes.fragment.TrackFragmentHeaderBox;
9import com.googlecode.mp4parser.AbstractFullBox;
10import com.googlecode.mp4parser.boxes.basemediaformat.TrackEncryptionBox;
11import com.googlecode.mp4parser.util.Path;
12
13import java.io.IOException;
14import java.math.BigInteger;
15import java.nio.ByteBuffer;
16import java.nio.channels.WritableByteChannel;
17import java.util.ArrayList;
18import java.util.Arrays;
19import java.util.LinkedList;
20import java.util.List;
21
22
23public abstract class AbstractSampleEncryptionBox extends AbstractFullBox {
24    int algorithmId = -1;
25    int ivSize = -1;
26    byte[] kid = new byte[]{-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1};
27    List<Entry> entries = new LinkedList<Entry>();
28
29    protected AbstractSampleEncryptionBox(String type) {
30        super(type);
31    }
32
33    public int getOffsetToFirstIV() {
34        int offset = (getSize() > (1l << 32) ? 16 : 8);
35        offset += isOverrideTrackEncryptionBoxParameters() ? 20 : 0;
36        offset += 4; //num entries
37        return offset;
38    }
39
40    @Override
41    public void _parseDetails(ByteBuffer content) {
42        parseVersionAndFlags(content);
43        int useThisIvSize = -1;
44        if ((getFlags() & 0x1) > 0) {
45            algorithmId = IsoTypeReader.readUInt24(content);
46            ivSize = IsoTypeReader.readUInt8(content);
47            useThisIvSize = ivSize;
48            kid = new byte[16];
49            content.get(kid);
50        } else {
51            List<Box> tkhds = Path.getPaths(this, "/moov[0]/trak/tkhd");
52            for (Box tkhd : tkhds) {
53                if (((TrackHeaderBox) tkhd).getTrackId() == this.getParent().getBoxes(TrackFragmentHeaderBox.class).get(0).getTrackId()) {
54                    AbstractTrackEncryptionBox tenc = (AbstractTrackEncryptionBox) Path.getPath(tkhd, "../mdia[0]/minf[0]/stbl[0]/stsd[0]/enc.[0]/sinf[0]/schi[0]/tenc[0]");
55                    if (tenc == null) {
56                        tenc = (AbstractTrackEncryptionBox) Path.getPath(tkhd, "../mdia[0]/minf[0]/stbl[0]/stsd[0]/enc.[0]/sinf[0]/schi[0]/uuid[0]");
57                    }
58                    useThisIvSize = tenc.getDefaultIvSize();
59                }
60            }
61        }
62        long numOfEntries = IsoTypeReader.readUInt32(content);
63
64        while (numOfEntries-- > 0) {
65            Entry e = new Entry();
66            e.iv = new byte[useThisIvSize < 0 ? 8 : useThisIvSize];  // default to 8
67            content.get(e.iv);
68            if ((getFlags() & 0x2) > 0) {
69                int numOfPairs = IsoTypeReader.readUInt16(content);
70                e.pairs = new LinkedList<Entry.Pair>();
71                while (numOfPairs-- > 0) {
72                    e.pairs.add(e.createPair(IsoTypeReader.readUInt16(content), IsoTypeReader.readUInt32(content)));
73                }
74            }
75            entries.add(e);
76
77        }
78    }
79
80
81    public int getSampleCount() {
82        return entries.size();
83    }
84
85    public List<Entry> getEntries() {
86        return entries;
87    }
88
89    public void setEntries(List<Entry> entries) {
90        this.entries = entries;
91    }
92
93    public int getAlgorithmId() {
94        return algorithmId;
95    }
96
97    public void setAlgorithmId(int algorithmId) {
98        this.algorithmId = algorithmId;
99    }
100
101    public int getIvSize() {
102        return ivSize;
103    }
104
105    public void setIvSize(int ivSize) {
106        this.ivSize = ivSize;
107    }
108
109    public byte[] getKid() {
110        return kid;
111    }
112
113    public void setKid(byte[] kid) {
114        this.kid = kid;
115    }
116
117
118    public boolean isSubSampleEncryption() {
119        return (getFlags() & 0x2) > 0;
120    }
121
122    public boolean isOverrideTrackEncryptionBoxParameters() {
123        return (getFlags() & 0x1) > 0;
124    }
125
126    public void setSubSampleEncryption(boolean b) {
127        if (b) {
128            setFlags(getFlags() | 0x2);
129        } else {
130            setFlags(getFlags() & (0xffffff ^ 0x2));
131        }
132    }
133
134    public void setOverrideTrackEncryptionBoxParameters(boolean b) {
135        if (b) {
136            setFlags(getFlags() | 0x1);
137        } else {
138            setFlags(getFlags() & (0xffffff ^ 0x1));
139        }
140    }
141
142
143    @Override
144    protected void getContent(ByteBuffer byteBuffer) {
145        writeVersionAndFlags(byteBuffer);
146        if (isOverrideTrackEncryptionBoxParameters()) {
147            IsoTypeWriter.writeUInt24(byteBuffer, algorithmId);
148            IsoTypeWriter.writeUInt8(byteBuffer, ivSize);
149            byteBuffer.put(kid);
150        }
151        IsoTypeWriter.writeUInt32(byteBuffer, entries.size());
152        for (Entry entry : entries) {
153            if (isOverrideTrackEncryptionBoxParameters()) {
154                byte[] ivFull = new byte[ivSize];
155                System.arraycopy(entry.iv, 0, ivFull, ivSize - entry.iv.length, entry.iv.length);
156                byteBuffer.put(ivFull);
157            } else {
158                // just put the iv - i don't know any better
159                byteBuffer.put(entry.iv);
160            }
161            if (isSubSampleEncryption()) {
162                IsoTypeWriter.writeUInt16(byteBuffer, entry.pairs.size());
163                for (Entry.Pair pair : entry.pairs) {
164                    IsoTypeWriter.writeUInt16(byteBuffer, pair.clear);
165                    IsoTypeWriter.writeUInt32(byteBuffer, pair.encrypted);
166                }
167            }
168        }
169    }
170
171    @Override
172    protected long getContentSize() {
173        long contentSize = 4;
174        if (isOverrideTrackEncryptionBoxParameters()) {
175            contentSize += 4;
176            contentSize += kid.length;
177        }
178        contentSize += 4;
179        for (Entry entry : entries) {
180            contentSize += entry.getSize();
181        }
182        return contentSize;
183    }
184
185    @Override
186    public void getBox(WritableByteChannel os) throws IOException {
187        super.getBox(os);
188    }
189
190    public Entry createEntry() {
191        return new Entry();
192    }
193
194    public class Entry {
195        public byte[] iv;
196        public List<Pair> pairs = new LinkedList<Pair>();
197
198        public int getSize() {
199            int size = 0;
200            if (isOverrideTrackEncryptionBoxParameters()) {
201                size = ivSize;
202            } else {
203                size = iv.length;
204            }
205
206
207            if (isSubSampleEncryption()) {
208                size += 2;
209                for (Entry.Pair pair : pairs) {
210                    size += 6;
211                }
212            }
213            return size;
214        }
215
216        public Pair createPair(int clear, long encrypted) {
217            return new Pair(clear, encrypted);
218        }
219
220
221        public class Pair {
222            public int clear;
223            public long encrypted;
224
225            public Pair(int clear, long encrypted) {
226                this.clear = clear;
227                this.encrypted = encrypted;
228            }
229
230            @Override
231            public boolean equals(Object o) {
232                if (this == o) {
233                    return true;
234                }
235                if (o == null || getClass() != o.getClass()) {
236                    return false;
237                }
238
239                Pair pair = (Pair) o;
240
241                if (clear != pair.clear) {
242                    return false;
243                }
244                if (encrypted != pair.encrypted) {
245                    return false;
246                }
247
248                return true;
249            }
250
251            @Override
252            public int hashCode() {
253                int result = clear;
254                result = 31 * result + (int) (encrypted ^ (encrypted >>> 32));
255                return result;
256            }
257
258            @Override
259            public String toString() {
260                return "clr:" + clear + " enc:" + encrypted;
261            }
262        }
263
264
265        @Override
266        public boolean equals(Object o) {
267            if (this == o) {
268                return true;
269            }
270            if (o == null || getClass() != o.getClass()) {
271                return false;
272            }
273
274            Entry entry = (Entry) o;
275
276            if (!new BigInteger(iv).equals(new BigInteger(entry.iv))) {
277                return false;
278            }
279            if (pairs != null ? !pairs.equals(entry.pairs) : entry.pairs != null) {
280                return false;
281            }
282
283            return true;
284        }
285
286        @Override
287        public int hashCode() {
288            int result = iv != null ? Arrays.hashCode(iv) : 0;
289            result = 31 * result + (pairs != null ? pairs.hashCode() : 0);
290            return result;
291        }
292
293        @Override
294        public String toString() {
295            return "Entry{" +
296                    "iv=" + Hex.encodeHex(iv) +
297                    ", pairs=" + pairs +
298                    '}';
299        }
300    }
301
302    @Override
303    public boolean equals(Object o) {
304        if (this == o) {
305            return true;
306        }
307        if (o == null || getClass() != o.getClass()) {
308            return false;
309        }
310
311        AbstractSampleEncryptionBox that = (AbstractSampleEncryptionBox) o;
312
313        if (algorithmId != that.algorithmId) {
314            return false;
315        }
316        if (ivSize != that.ivSize) {
317            return false;
318        }
319        if (entries != null ? !entries.equals(that.entries) : that.entries != null) {
320            return false;
321        }
322        if (!Arrays.equals(kid, that.kid)) {
323            return false;
324        }
325
326        return true;
327    }
328
329    @Override
330    public int hashCode() {
331        int result = algorithmId;
332        result = 31 * result + ivSize;
333        result = 31 * result + (kid != null ? Arrays.hashCode(kid) : 0);
334        result = 31 * result + (entries != null ? entries.hashCode() : 0);
335        return result;
336    }
337
338    public List<Short> getEntrySizes() {
339        List<Short> entrySizes = new ArrayList<Short>(entries.size());
340        for (Entry entry : entries) {
341            short size = (short) entry.iv.length;
342            if (isSubSampleEncryption()) {
343                size += 2; //numPairs
344                size += entry.pairs.size() * 6;
345            }
346            entrySizes.add(size);
347        }
348        return entrySizes;
349    }
350}
351