|
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