Initial revision
[enigma2-plugins.git] / reconstructapsc / src_cc / reconstruct_apsc.cc
1  /* Copyright (C) 2009 Anders Holst
2   *
3   * This program is free software; you can redistribute it and/or
4   * modify it under the terms of the GNU General Public License as
5   * published by the Free Software Foundation; either version 2, or
6   * (at your option) any later version.
7   * 
8   * This program is distributed in the hope that it will be useful,
9   * but WITHOUT ANY WARRANTY; without even the implied warranty of
10   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11   * GNU General Public License for more details.
12   */
13
14 #define _LARGEFILE64_SOURCE
15 #include <sys/types.h>
16 #include <sys/stat.h>
17 #include <fcntl.h>
18 #include <dirent.h>
19 #include <unistd.h>
20 #include <stdio.h>
21 #include <stdlib.h>
22 #include <string.h>
23 #include <byteswap.h>
24 #include <errno.h>
25
26 #define LEN 24064
27
28
29 char* makefilename(const char* dir, const char* base, const char* ext, const char* post)
30 {
31   static char buf[256];
32   int len1, len2, len3;
33   len1 = (dir ? strlen(dir) : 0);
34   len2 = (base ? strlen(base) : 0);
35   len3 = (ext ? strlen(ext) : 0);
36   if (dir) {
37     strcpy(buf, dir);
38     if (buf[len1-1] != '/') {
39       buf[len1++] = '/';
40       buf[len1] = 0;
41     }
42   }
43   if (base)
44     strcpy(buf+len1, base);
45   if (ext && len2>=len3 && !strcmp(base+len2-len3,ext))
46     len2 -= len3;
47   if (ext)
48     strcpy(buf+len1+len2, ext);
49   if (post)
50     strcpy(buf+len1+len2+len3, post);
51   return buf;
52 }
53
54 int writebufinternal(int f, off64_t sz, off64_t tm)
55 {
56   off64_t buf[2];
57   buf[0] = (off64_t)bswap_64((unsigned long long int)sz);
58   buf[1] = (off64_t)bswap_64((unsigned long long int)tm);
59   if (write(f, buf, 16) != 16)
60     return 1;
61   else
62     return 0;
63 }
64
65 int framepid(unsigned char* buf, int pos)
66 {
67   return ((buf[pos+1] & 0x1f) << 8) + buf[pos+2];
68 }
69
70 off64_t framepts(unsigned char* buf, int pos)
71 {
72   int tmp = (buf[pos+3] & 0x20 ? pos+buf[pos+4]+5 : pos+4);
73   off64_t pts;
74   if (buf[pos+1] & 0x40 &&
75       buf[pos+3] & 0x10 &&
76       buf[tmp]==0 && buf[tmp+1]==0 && buf[tmp+2]==1 &&
77       buf[tmp+7] & 0x80) {
78     pts  = ((unsigned long long)(buf[tmp+9]&0xE))  << 29;
79     pts |= ((unsigned long long)(buf[tmp+10]&0xFF)) << 22;
80     pts |= ((unsigned long long)(buf[tmp+11]&0xFE)) << 14;
81     pts |= ((unsigned long long)(buf[tmp+12]&0xFF)) << 7;
82     pts |= ((unsigned long long)(buf[tmp+13]&0xFE)) >> 1;
83   } else
84     pts = 0;
85   return pts;
86 }
87
88 int framesearch(int fts, int first, off64_t& retpos, off64_t& retpts, off64_t& retpos2, off64_t& retdat)
89 {
90   static unsigned char buf[LEN];
91   static int ind;
92   static off64_t pos = -1;
93   static off64_t num;
94   static int pid = -1;
95   static int st = 0;
96   static int sdflag = 0;
97   unsigned char* p;
98   if (pos == -1 || first) {
99     num = read(fts, buf, LEN);
100     ind = 0;
101     pos = 0;
102     st = 0;
103   }
104   while (1) {
105     p = buf+ind+st;
106     ind = -1;
107     for (; p < buf+num-6; p++)
108       if (p[0]==0 && p[1]==0 && p[2]==1) {
109         ind = ((p - buf)/188)*188;
110         if ((p[3] & 0xf0) == 0xe0 && (buf[ind+1] & 0x40) &&
111             (p-buf)-ind == (buf[ind+3] & 0x20 ? buf[ind+4] + 5 : 4)) {
112           pid = framepid(buf, ind);
113         } else if (pid != -1 && pid != framepid(buf, ind)) {
114           ind = -1;
115           continue;
116         }
117         if (p[3]==0 || p[3]==0xb3 || p[3]==0xb8) { // MPEG2
118           if (p[3]==0 && ((p[5] >> 3 & 7) == 1)) {
119             retpts = framepts(buf, ind);
120             retpos = pos + ind;
121           } else {
122             retpts = 0;
123             retpos = -1;
124           }
125           retdat = (unsigned int) p[3] | (p[4]<<8) | (p[5]<<16) | (p[6]<<24);
126           retpos2 = pos + (p - buf);
127           st = (p - buf) - ind + 1;
128           sdflag = 1;
129           return 1; 
130         } else if (!sdflag && p[3]==0x09 && (buf[ind+1] & 0x40)) { // H264
131           if ((p[4] >> 5)==0) {
132             retpts = framepts(buf, ind);
133             retpos = pos + ind;
134           } else {
135             retpts = 0;
136             retpos = -1;
137           }
138           retdat = p[3] | (p[4]<<8);
139           retpos2 = pos + (p - buf);
140           st = (p - buf) - ind + 1;
141           return 1; 
142         } else {
143           ind = -1;
144           continue;
145         }
146       }
147     st = 0;
148     sdflag = 0; // reset to get some fault tolerance
149     if (num == LEN) {
150       pos += num;
151       num = read(fts, buf, LEN);
152       ind = 0;
153     } else if (num) {
154       ind = num;
155       retpts = 0;
156       retdat = 0;
157       retpos = pos + num;
158       num = 0;
159       return -1;
160     } else {
161       retpts = 0;
162       retdat = 0;
163       retpos = 0;
164       return -1;
165     }
166   }
167 }
168
169 int do_one(int fts, int fap, int fsc)
170 {
171   off64_t pos;
172   off64_t pos2;
173   off64_t pts;
174   off64_t dat;
175   int first = 1;
176   while (framesearch(fts, first, pos, pts, pos2, dat) >= 0) {
177     first = 0;
178     if (pos >= 0)
179       if (fap >= 0 && writebufinternal(fap, pos, pts))
180         return 1;
181     if (fsc >= 0 && writebufinternal(fsc, pos2, dat))
182       return 1;
183   }
184   return 0;
185 }
186
187 int do_movie(char* inname)
188 {
189   int f_ts=-1, f_sc=-1, f_ap=-1, f_tmp=-1;
190   char* tmpname;
191   tmpname = makefilename(0, inname, ".ts", 0);
192   f_ts = open(tmpname, O_RDONLY | O_LARGEFILE);
193   if (f_ts == -1) {
194     printf("Failed to open input stream file \"%s\"\n", tmpname);
195     return 1;
196   }
197   tmpname = makefilename(0, inname, ".ts", ".reconstruct_apsc");
198   f_tmp = open(tmpname, O_WRONLY | O_CREAT | O_TRUNC, 0x1a4);
199   if (f_tmp == -1) {
200     printf("Failed to open sentry file \"%s\"\n", tmpname);
201     goto failure;
202   }
203   close(f_tmp);
204   tmpname = makefilename(0, inname, ".ts", ".ap");
205   f_ap = open(tmpname, O_WRONLY | O_CREAT | O_TRUNC, 0x1a4);
206   if (f_ap == -1) {
207     printf("Failed to open output .ap file \"%s\"\n", tmpname);
208     goto failure;
209   }
210   tmpname = makefilename(0, inname, ".ts", ".sc");
211   f_sc = open(tmpname, O_WRONLY | O_CREAT | O_TRUNC, 0x1a4);
212   if (f_sc == -1) {
213     printf("Failed to open output .sc file \"%s\"\n", tmpname);
214     goto failure;
215   }
216
217   printf("  Processing .ap and .sc of \"%s\" ... ", inname);
218   fflush(stdout);
219   if (do_one(f_ts, f_ap, f_sc)) {
220     printf("\nFailed to reconstruct files for \"%s\"\n", inname);
221     goto failure;
222   }
223   printf("done\n");
224
225   close(f_ts);
226   close(f_ap);
227   close(f_sc);
228   unlink(makefilename(0, inname, ".ts", ".reconstruct_apsc"));
229   return 0;
230  failure:
231   if (f_ts != -1)
232     close(f_ts);
233   if (f_ap != -1) {
234     close(f_ap);
235     unlink(makefilename(0, inname, ".ts", ".ap"));
236   }
237   if (f_sc != -1) {
238     close(f_sc);
239     unlink(makefilename(0, inname, ".ts", ".sc"));
240   }
241   unlink(makefilename(0, inname, ".ts", ".reconstruct_apsc"));
242   return 1;
243 }
244
245 int do_directory(char* dirname)
246 {
247   int f_ts, f_sc, f_ap, f_tmp;
248   int do_ap, do_sc;
249   char *inname, *tmpname;
250   DIR* dir = opendir(dirname);
251   dirent* entry;
252   struct stat statbuf;
253   if (dir) {
254     while ((entry = readdir(dir))) {
255       inname = entry->d_name;
256       if (strlen(inname) > 3 && !strcmp(inname + strlen(inname) - 3, ".ts")) {
257         tmpname = makefilename(dirname, inname, ".ts", ".reconstruct_apsc");
258         errno = 0;
259         if (stat(tmpname, &statbuf) != -1)
260           do_ap = do_sc = 1;
261         else {
262           tmpname = makefilename(dirname, inname, ".ts", ".ap");
263           errno = 0;
264           do_ap = (stat(tmpname, &statbuf) == -1 && errno == ENOENT);
265           tmpname = makefilename(dirname, inname, ".ts", ".sc");
266           errno = 0;
267           do_sc = (stat(tmpname, &statbuf) == -1 && errno == ENOENT);
268         }
269         if (do_ap || do_sc) {
270           f_ts=-1, f_sc=-1, f_ap=-1, f_tmp=-1;
271           tmpname = makefilename(dirname, inname, ".ts", 0);
272           f_ts = open(tmpname, O_RDONLY | O_LARGEFILE);
273           if (f_ts == -1) {
274             printf("Failed to open input stream file \"%s\"\n", tmpname);
275             continue;
276           }
277           tmpname = makefilename(dirname, inname, ".ts", ".reconstruct_apsc");
278           f_tmp = open(tmpname, O_WRONLY | O_CREAT | O_TRUNC, 0x1a4);
279           if (f_tmp == -1) {
280             printf("Failed to open sentry file \"%s\"\n", tmpname);
281             goto failure;
282           }
283           close(f_tmp);
284           if (do_ap) {
285             tmpname = makefilename(dirname, inname, ".ts", ".ap");
286             f_ap = open(tmpname, O_WRONLY | O_CREAT | O_TRUNC, 0x1a4);
287             if (f_ap == -1) {
288               printf("Failed to open output .ap file \"%s\"\n", tmpname);
289               goto failure;
290             }
291           } else
292             f_ap = -1;
293           if (do_sc) {
294             tmpname = makefilename(dirname, inname, ".ts", ".sc");
295             f_sc = open(tmpname, O_WRONLY | O_CREAT | O_TRUNC, 0x1a4);
296             if (f_sc == -1) {
297               printf("Failed to open output .sc file \"%s\"\n", tmpname);
298               goto failure;
299             }
300           } else
301             f_sc = -1;
302
303           printf("  Processing %s of \"%s\" ... ", (do_ap ? (do_sc ? ".ap and .sc" : ".ap") : ".sc"), inname);
304           fflush(stdout);
305           if (do_one(f_ts, f_ap, f_sc)) {
306             printf("\nFailed to reconstruct files for \"%s\"\n", inname);
307             close(f_ts);
308             if (f_ap != -1) {
309               close(f_ap);
310               unlink(makefilename(dirname, inname, ".ts", ".ap"));
311             }
312             if (f_sc != -1) {
313               close(f_sc);
314               unlink(makefilename(dirname, inname, ".ts", ".sc"));
315             }
316             unlink(makefilename(dirname, inname, ".ts", ".reconstruct_apsc"));
317             continue;
318           }
319           printf("done\n");
320
321           close(f_ts);
322           close(f_ap);
323           close(f_sc);
324           unlink(makefilename(dirname, inname, ".ts", ".reconstruct_apsc"));
325         }
326       }
327     }
328     closedir(dir);
329   } else {
330     printf("Failed to open directory \"%s\"\n", dirname);
331     return 1;
332   }
333   return 0;
334  failure:
335   closedir(dir);
336   if (f_ts != -1)
337     close(f_ts);
338   if (f_ap != -1) {
339     close(f_ap);
340     unlink(makefilename(dirname, inname, ".ts", ".ap"));
341   }
342   if (f_sc != -1) {
343     close(f_sc);
344     unlink(makefilename(dirname, inname, ".ts", ".sc"));
345   }
346   unlink(makefilename(dirname, inname, ".ts", ".reconstruct_apsc"));
347   return 1;
348 }
349
350 int main(int argc, char* argv[])
351 {
352   if (argc == 2 && *argv[1] != '-') {
353     if (do_movie(argv[1]))
354       exit(1);
355   } else if (argc == 3 && !strcmp(argv[1], "-d")) {
356     if (do_directory(argv[2]))
357       exit(1);
358   } else {
359     printf("Usage: reconstruct_apsc movie_file\n");
360     printf("       reconstruct_apsc -d movie_directory\n");
361     exit(1);
362   }
363 }
364