|
17 | 17 | from .usm import is_usm, Usm, Vp9, H264, OpMode, generate_keys |
18 | 18 |
|
19 | 19 |
|
| 20 | +def create_usm(): |
| 21 | + parser = argparse.ArgumentParser("WannaCRI Create USM", allow_abbrev=False) |
| 22 | + parser.add_argument( |
| 23 | + "operation", |
| 24 | + metavar="operation", |
| 25 | + type=str, |
| 26 | + choices=OP_LIST, |
| 27 | + help="Specify operation.", |
| 28 | + ) |
| 29 | + parser.add_argument( |
| 30 | + "input", |
| 31 | + metavar="input file path", |
| 32 | + type=existing_file, |
| 33 | + help="Path to video file.", |
| 34 | + ) |
| 35 | + parser.add_argument( |
| 36 | + "-e", |
| 37 | + "--encoding", |
| 38 | + type=str, |
| 39 | + default="shift-jis", |
| 40 | + help="Character encoding used in creating USM. Defaults to shift-jis.", |
| 41 | + ) |
| 42 | + parser.add_argument( |
| 43 | + "-o", |
| 44 | + "--output", |
| 45 | + type=dir_path, |
| 46 | + default=None, |
| 47 | + help="Output path. Defaults to the same place as input.", |
| 48 | + ) |
| 49 | + parser.add_argument( |
| 50 | + "--ffprobe", |
| 51 | + type=str, |
| 52 | + default=".", |
| 53 | + help="Path to ffprobe executable or directory. Defaults to CWD.", |
| 54 | + ) |
| 55 | + parser.add_argument( |
| 56 | + "-k", "--key", type=key, default=None, help="Encryption key for encrypted USMs." |
| 57 | + ) |
| 58 | + args = parser.parse_args() |
| 59 | + |
| 60 | + ffprobe_path = find_ffprobe(args.ffprobe) |
| 61 | + |
| 62 | + # TODO: Add support for more video codecs and audio codecs |
| 63 | + codec = Sofdec2Codec.from_file(args.input) |
| 64 | + if codec is Sofdec2Codec.VP9: |
| 65 | + video = Vp9(args.input, ffprobe_path=ffprobe_path) |
| 66 | + elif codec is Sofdec2Codec.H264: |
| 67 | + video = H264(args.input, ffprobe_path=ffprobe_path) |
| 68 | + else: |
| 69 | + raise NotImplementedError("Non-Vp9/H.264 files are not yet implemented.") |
| 70 | + |
| 71 | + filename = os.path.splitext(args.input)[0] |
| 72 | + |
| 73 | + usm = Usm(videos=[video], key=args.key) |
| 74 | + with open(filename + ".usm", "wb") as f: |
| 75 | + mode = OpMode.NONE if args.key is None else OpMode.ENCRYPT |
| 76 | + |
| 77 | + for packet in usm.stream(mode, encoding=args.encoding): |
| 78 | + f.write(packet) |
| 79 | + |
| 80 | + print("Done creating USM file.") |
| 81 | + |
| 82 | + |
20 | 83 | def extract_usm(): |
21 | 84 | """One of the main functions in the command-line program. Extracts a USM or extracts |
22 | 85 | multiple USMs given a path as input.""" |
@@ -242,69 +305,6 @@ def probe_usm(): |
242 | 305 | print(f'Probe complete. All logs are stored in "{args.output}" folder') |
243 | 306 |
|
244 | 307 |
|
245 | | -def create_usm(): |
246 | | - parser = argparse.ArgumentParser("WannaCRI Create USM", allow_abbrev=False) |
247 | | - parser.add_argument( |
248 | | - "operation", |
249 | | - metavar="operation", |
250 | | - type=str, |
251 | | - choices=OP_LIST, |
252 | | - help="Specify operation.", |
253 | | - ) |
254 | | - parser.add_argument( |
255 | | - "input", |
256 | | - metavar="input file path", |
257 | | - type=existing_file, |
258 | | - help="Path to video file.", |
259 | | - ) |
260 | | - parser.add_argument( |
261 | | - "-e", |
262 | | - "--encoding", |
263 | | - type=str, |
264 | | - default="shift-jis", |
265 | | - help="Character encoding used in creating USM. Defaults to shift-jis.", |
266 | | - ) |
267 | | - parser.add_argument( |
268 | | - "-o", |
269 | | - "--output", |
270 | | - type=dir_path, |
271 | | - default=None, |
272 | | - help="Output path. Defaults to the same place as input.", |
273 | | - ) |
274 | | - parser.add_argument( |
275 | | - "--ffprobe", |
276 | | - type=str, |
277 | | - default=".", |
278 | | - help="Path to ffprobe executable or directory. Defaults to CWD.", |
279 | | - ) |
280 | | - parser.add_argument( |
281 | | - "-k", "--key", type=key, default=None, help="Encryption key for encrypted USMs." |
282 | | - ) |
283 | | - args = parser.parse_args() |
284 | | - |
285 | | - ffprobe_path = find_ffprobe(args.ffprobe) |
286 | | - |
287 | | - # TODO: Add support for more video codecs and audio codecs |
288 | | - codec = Sofdec2Codec.from_file(args.input) |
289 | | - if codec is Sofdec2Codec.VP9 or codec is Sofdec2Codec.H264: |
290 | | - if codec is Sofdec2Codec.VP9: |
291 | | - video = Vp9(args.input, ffprobe_path=ffprobe_path) |
292 | | - elif codec is Sofdec2Codec.H264: |
293 | | - video = H264(args.input, ffprobe_path=ffprobe_path) |
294 | | - filename = os.path.splitext(args.input)[0] |
295 | | - |
296 | | - usm = Usm(videos=[video], key=args.key) |
297 | | - with open(filename + ".usm", "wb") as f: |
298 | | - mode = OpMode.NONE if args.key is None else OpMode.ENCRYPT |
299 | | - |
300 | | - for packet in usm.stream(mode, encoding=args.encoding): |
301 | | - f.write(packet) |
302 | | - else: |
303 | | - raise NotImplementedError("Non-Vp9/H.264 files are not yet implemented.") |
304 | | - |
305 | | - print("Done creating USM file.") |
306 | | - |
307 | | - |
308 | 308 | def encrypt_usm(): |
309 | 309 | parser = argparse.ArgumentParser("WannaCRI Encrypt USM/s", allow_abbrev=False) |
310 | 310 | parser.add_argument( |
|
0 commit comments