add ds mov uploading?

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.

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

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: