diff --git a/theunderground/forms.py b/theunderground/forms.py index 33340001..a9f65892 100644 --- a/theunderground/forms.py +++ b/theunderground/forms.py @@ -72,6 +72,7 @@ def validate_new_password(self, _): class MovieUploadForm(FlaskForm): movie = FileField("Movie", validators=[FileRequired()]) + dsmovie = FileField("DS Movie") title = StringField("Movie title", validators=[DataRequired(), Length(max=48)]) thumbnail = FileField("Movie thumbnail", validators=[FileRequired()]) # Choices for the select field are only evaluated once, so we must set it when necessary. diff --git a/theunderground/mobiclip.py b/theunderground/mobiclip.py index 5e04d24b..6fc113df 100644 --- a/theunderground/mobiclip.py +++ b/theunderground/mobiclip.py @@ -36,6 +36,13 @@ def get_movie_path(movie_id: int) -> str: return f"movie/{movie_byte}" else: return f"./assets/movies/{movie_byte}" + +def get_ds_movie_path(movie_id: int) -> str: + movie_byte = get_movie_byte(movie_id) + if s3: + return f"dsmov/{movie_byte}" + else: + return f"./assets/dsmov/{movie_byte}" def get_pay_movie_dir(movie_id: int) -> str: @@ -54,6 +61,19 @@ def validate_mobiclip(file_data: bytes) -> bool: return True +def validate_mobi_dsi(file_data: bytes) -> bool: + # Validate file magic + if file_data[0:4] == b"MODS": + pass + elif file_data[0:4] == b"SSCF": + # Check for NINTENDO header + if file_data[17:25] != b"NINTENDO": + return False + else: + return False + + return True + def get_category_list(): db_categories = Categories.query.all() @@ -92,8 +112,9 @@ def get_mobiclip_length(file_data: bytes) -> str: return strftime("%H:%M:%S", gmtime(length)) -def save_movie_data(movie_id: int, thumbnail_data: bytes, movie_data: bytes): +def save_movie_data(movie_id: int, thumbnail_data: bytes, movie_data: bytes, ds_movie_data: bytes = None): movie_dir = get_movie_path(movie_id) + ds_movie_dir = get_ds_movie_path(movie_id) md5_hash = get_movie_byte(movie_id) if s3: @@ -112,6 +133,11 @@ def save_movie_data(movie_id: int, thumbnail_data: bytes, movie_data: bytes): s3.upload_fileobj( BytesIO(thumbnail_data), config.r2_bucket_name, thumbnail_path ) + + if ds_movie_data: + # Upload DS movie + ds_movie_path = f"{ds_movie_dir}/{movie_id}.enc" + s3.upload_fileobj(BytesIO(ds_movie_data), config.r2_bucket_name, ds_movie_path) else: if not os.path.isdir(movie_dir): os.makedirs(movie_dir) @@ -127,6 +153,15 @@ def save_movie_data(movie_id: int, thumbnail_data: bytes, movie_data: bytes): movie.write(movie_data) movie.close() + if ds_movie_data: + if not os.path.isdir(ds_movie_dir): + os.makedirs(ds_movie_dir) + + # Write DS movie + ds_movie = open(f"{ds_movie_dir}/{movie_id}.enc", "wb") + ds_movie.write(ds_movie_data) + ds_movie.close() + def save_pay_movie_data( movie_id: int, thumbnail_data: bytes, movie_data: bytes, poster_data: bytes diff --git a/theunderground/movies.py b/theunderground/movies.py index 1f88217a..7d5834df 100644 --- a/theunderground/movies.py +++ b/theunderground/movies.py @@ -13,6 +13,7 @@ from theunderground.mobiclip import ( get_category_list, validate_mobiclip, + validate_mobi_dsi, get_mobiclip_length, save_movie_data, delete_movie_data, @@ -54,11 +55,21 @@ def add_movie(): if form.validate_on_submit(): movie = form.movie.data + dsmovie = form.dsmovie.data thumbnail = form.thumbnail.data if movie and thumbnail: movie_data = movie[0].read() thumbnail_data = thumbnail[0].read() + if dsmovie: + dsmovie_data = dsmovie[0].read() + genre = 1 + validation_ds = validate_mobi_dsi(dsmovie_data) + else: + dsmovie_data = None + genre = 0 + validation_ds = False + if validate_mobiclip(movie_data): # Get the Mobiclip's length from header. length = get_mobiclip_length(movie_data) @@ -70,16 +81,19 @@ def add_movie(): category_id=form.category.data, length=length, aspect=True, - genre=0, + genre=genre, sp_page_id=0, staff=False, ) + if dsmovie and validation_ds: + db_movie.ds_mov_id = db_movie.movie_id + db.session.add(db_movie) db.session.commit() # Now that we've inserted the movie, we can properly move it. - save_movie_data(db_movie.movie_id, thumbnail_data, movie_data) + save_movie_data(db_movie.movie_id, thumbnail_data, movie_data, dsmovie_data) # Finally update the category if needed by S3 if s3: