diff --git a/main.py b/main.py index ebc8f36..b87da07 100644 --- a/main.py +++ b/main.py @@ -39,6 +39,7 @@ import requests import json +import tempfile # --- --- --- Section Init --- --- --- # # Login to RainClassroom @@ -49,6 +50,10 @@ DOWNLOAD_FOLDER = "data" CACHE_FOLDER = "cache" +# Make a instance-specific cache folder +TEMP_FOLDER = tempfile.mkdtemp() +print(f"Temp Folder: {TEMP_FOLDER}") + os.makedirs(DOWNLOAD_FOLDER, exist_ok=True) os.makedirs(CACHE_FOLDER, exist_ok=True) @@ -217,6 +222,7 @@ if os.path.exists(f"{DOWNLOAD_FOLDER}/{name_prefix}.mp4"): print(f"Skipping {name_prefix} - Video already present") + time.sleep(0.5) return has_error = False @@ -233,15 +239,25 @@ if not has_error and len(lesson_video_data['data']['live']) > 0: print(f"Concatenating {name_prefix}") - with open(f"{CACHE_FOLDER}/concat.txt", "w") as f: + ffmpeg_input_file = f"{TEMP_FOLDER}/concat.txt" + + # Get absolute path of the video files + cache_absolute = os.path.abspath(f"{CACHE_FOLDER}") + + with open(ffmpeg_input_file, "w") as f: f.write("\n".join( - [f"file '{name_prefix}-{i}.mp4'" for i in range(len(lesson_video_data['data']['live']))] + [f"file '{cache_absolute}/{name_prefix}-{i}.mp4'" for i in range(len(lesson_video_data['data']['live']))] )) - cmd = f"ffmpeg -f concat -safe 0 -hwaccel cuda -hwaccel_output_format cuda -i {CACHE_FOLDER}/concat.txt -c:v hevc_nvenc -b:v 200k -maxrate 400k -bufsize 3200k -r 8 -rc-lookahead 1024 -c:a copy -rematrix_maxval 1.0 -ac 1 '{DOWNLOAD_FOLDER}/{name_prefix}.mp4' -n" + cmd = f"ffmpeg -f concat -safe 0 -hwaccel cuda -hwaccel_output_format cuda -i {ffmpeg_input_file} -c:v hevc_nvenc -b:v 200k -maxrate 400k -bufsize 3200k -r 8 -rc-lookahead 1024 -c:a copy -rematrix_maxval 1.0 -ac 1 '{DOWNLOAD_FOLDER}/{name_prefix}.mp4' -n" print(cmd) - os.system(cmd) + ret = os.system(cmd) + if ret != 0: + if ret == 130: + raise KeyboardInterrupt("Interrupted") + + raise Exception(f"Failed to concatenate {name_prefix}") if has_error: with open(f"{DOWNLOAD_FOLDER}/error.log", "a") as f: @@ -272,6 +288,9 @@ ret = os.system( f"aria2c -o '{CACHE_FOLDER}/{name_prefix}-{order}.mp4' -x 4 -s 4 '{url}' -c") if ret != 0: + if ret == 130: + raise KeyboardInterrupt("Interrupted") + raise Exception(f"Failed to download {name_prefix}-{order}") # --- --- --- Section Download Lesson PPT --- --- --- # @@ -438,13 +457,16 @@ # If PDF is present, skip if os.path.exists(f"{DOWNLOAD_FOLDER}/{name_prefix}.pdf"): print(f"Skipping {name_prefix} - PDF already present") + time.sleep(0.5) return os.makedirs(f"{DOWNLOAD_FOLDER}/{name_prefix}", exist_ok=True) images = [] - with open(f"{CACHE_FOLDER}/ppt_download.txt", "w") as f: + aria2_input_file = f"{TEMP_FOLDER}/ppt_download.txt" + + with open(aria2_input_file, "w") as f: for slide in ppt_raw_data['data']['slides']: if not slide.get('cover'): continue @@ -459,7 +481,11 @@ # with open(f"{DOWNLOAD_FOLDER}/{name_prefix}/{slide['index']}.jpg", "wb") as f: # f.write(requests.get(slide['cover']).content) - os.system(f"aria2c -i {CACHE_FOLDER}/ppt_download.txt -x 16 -s 16 -c") + ret = os.system(f"aria2c -i {aria2_input_file} -x 16 -s 16 -c") + if ret != 0: + if ret == 7: + raise KeyboardInterrupt("Interrupted") + raise Exception(f"Failed to download {name_prefix}") from PIL import Image @@ -513,6 +539,16 @@ for course in courses: try: get_lesson_list(course) + except KeyboardInterrupt: + # Remove temp folder + print("Removing Temp Folder") + os.system(f"rm -rf {TEMP_FOLDER}") + print("Interrupted") + exit() except Exception as e: print(e) print(f"Failed to parse {course['name']}", file=sys.stderr) + +# Remove temp folder +print("Removing Temp Folder") +os.system(f"rm -rf {TEMP_FOLDER}") \ No newline at end of file