Add support for S3

asset.py

@@ -1,11 +1,14 @@
 import io
+import os
 import pathlib
 from typing import Union
 
 import flask
 from PIL import ImageFile, Image
-from flask import send_from_directory
+from flask import send_from_directory, redirect
 from wtforms import FileField
+from room import s3
+import config
 
 
 class Asset(object):
@@ -20,10 +23,17 @@ class Asset(object):
     # The filename of this asset within its directory.
     asset_name: str
 
+    s3_tag: str
+
+    always_save = False
+
     def asset_path(self) -> pathlib.Path:
         """The full path to this asset."""
         return self.asset_dir / self.asset_name
 
+    def s3_path(self) -> str:
+        return self.s3_tag + "/" + self.asset_name
+
     def ensure_exists(self):
         """Creates the parent directories for this asset."""
         self.asset_dir.mkdir(parents=True, exist_ok=True)
@@ -50,9 +60,37 @@ def encode(self, in_bytes: Union[bytes, FileField]):
 
         im = im.resize(self.dimensions)
 
-        # These defaults are required for the Wii to read an JPEG.
-        im.save(self.asset_path(), "jpeg", subsampling="4:2:0", progressive=False)
+        if s3:
+            img = io.BytesIO()
+            im.save(img, "jpeg", subsampling="4:2:0", progressive=False)
+            self.upload_to_s3(img.getvalue())
+
+            if self.always_save:
+                with open(self.asset_path(), "wb") as f:
+                    f.write(img.getvalue())
+        else:
+            # These defaults are required for the Wii to read an JPEG.
+            im.save(self.asset_path(), "jpeg", subsampling="4:2:0", progressive=False)
+
+    def delete(self):
+        if s3:
+            self.remove_from_s3()
+        else:
+            os.remove(self.asset_path())
 
     def send_file(self) -> flask.Response:
         """Wraps around Flask's send_from_directory method."""
         return send_from_directory(self.asset_dir, self.asset_name)
+
+    def upload_to_s3(self, im: bytes):
+        # I don't know why we have to convert the BytesIO object to bytes then back, but if we don't it will upload
+        # nothing.
+        s3.upload_fileobj(
+            io.BytesIO(im),
+            config.r2_bucket_name,
+            self.s3_path(),
+            ExtraArgs={"ContentType": "image/jpeg"},
+        )
+
+    def remove_from_s3(self):
+        s3.delete_object(Bucket=config.r2_bucket_name, Key=self.s3_path())

asset_data.py

@@ -1,8 +1,15 @@
+import io
+
+import config
+
 from asset import Asset
 from typing import Union
 from wtforms import FileField
+from room import s3
+from werkzeug.exceptions import NotFound
 
-import os
+from url1.wall_metadata import wall_metadata
+from url1.category_n import list_category_n
 
 
 class RoomLogoAsset(Asset):
@@ -12,6 +19,16 @@ def __init__(self, room_id: int):
         self.dimensions = (320, 180)
         self.asset_dir = self.base_asset_dir / "special" / f"{room_id}"
         self.asset_name = "f1234.img"
+        self.s3_tag = f"special/{room_id}/img"
+
+    def upload_to_s3(self, im: bytes):
+        # First the image
+        s3.upload_fileobj(
+            io.BytesIO(im),
+            config.r2_bucket_name,
+            self.s3_path(),
+            ExtraArgs={"ContentType": "image/jpeg"},
+        )
 
 
 class ParadeBannerAsset(Asset):
@@ -21,6 +38,71 @@ def __init__(self, room_id: int):
         self.dimensions = (184, 80)
         self.asset_dir = self.base_asset_dir / "special" / f"{room_id}"
         self.asset_name = "parade_banner.jpg"
+        self.room_id = room_id
+        self.always_save = True
+
+    def _upload_xmls_to_s3(self):
+        from url1.special.allbin import special_allbin
+
+        allbin_xml = special_allbin()
+        s3.upload_fileobj(
+            io.BytesIO(allbin_xml),
+            config.r2_bucket_name,
+            "special/allbin.xml",
+            ExtraArgs={"ContentType": "text/xml"},
+        )
+
+        from url1.special.all import special_all
+
+        all_xml = special_all()
+        s3.upload_fileobj(
+            io.BytesIO(all_xml),
+            config.r2_bucket_name,
+            "special/all.xml",
+            ExtraArgs={"ContentType": "text/xml"},
+        )
+
+        from url1.special.contact import special_contact_n
+
+        contact_xml = special_contact_n(self.room_id)
+
+        if not isinstance(contact_xml, NotFound):
+            s3.upload_fileobj(
+                io.BytesIO(contact_xml),
+                config.r2_bucket_name,
+                f"special/{self.room_id}/contact.xml",
+                ExtraArgs={"ContentType": "text/xml"},
+            )
+
+        from url1.special.page import special_page_n
+
+        page_xml = special_page_n(self.room_id)
+        if not isinstance(contact_xml, NotFound):
+            s3.upload_fileobj(
+                io.BytesIO(page_xml),
+                config.r2_bucket_name,
+                f"special/{self.room_id}/page.xml",
+                ExtraArgs={"ContentType": "text/xml"},
+            )
+
+    def upload_to_s3(self, im: bytes):
+        # For this specifically we need to upload to both S3 and save to disk.
+        s3.upload_fileobj(
+            io.BytesIO(im),
+            config.r2_bucket_name,
+            f"special/{self.room_id}/img/g1234.img",
+            ExtraArgs={"ContentType": "image/jpeg"},
+        )
+
+        self._upload_xmls_to_s3()
+
+    def remove_from_s3(self):
+        s3.delete_object(
+            Bucket=config.r2_bucket_name, Key=f"special/{self.room_id}/img/g1234.img"
+        )
+
+        # Update the XMLs
+        self._upload_xmls_to_s3()
 
 
 class NormalCategoryAsset(Asset):
@@ -31,6 +113,26 @@ def __init__(self, category_id: int):
         self.asset_dir = self.base_asset_dir / "normal-category"
         self.asset_name = f"{category_id}.img"
 
+        self.raw_asset_name = category_id
+        self.s3_tag = "list/category/img"
+
+    def upload_to_s3(self, im: bytes):
+        # First the image
+        s3.upload_fileobj(
+            io.BytesIO(im),
+            config.r2_bucket_name,
+            self.s3_path(),
+            ExtraArgs={"ContentType": "image/jpeg"},
+        )
+
+        # Next mirror the updates to the categories xml
+        s3.upload_fileobj(
+            io.BytesIO(list_category_n("01")),
+            config.r2_bucket_name,
+            "list/category/01.xml",
+            ExtraArgs={"ContentType": "text/xml"},
+        )
+
 
 class PayCategoryAsset(Asset):
     """Used for categories within the Theater."""
@@ -86,9 +188,29 @@ def __init__(self, seq: int, is_theatre: bool):
         self.dimensions = (256, 360)
         self.asset_dir = self.base_asset_dir / asset_dir
         self.asset_name = asset_name
+        self.raw_asset_name = seq
 
-    def delete(self):
-        os.remove(self.asset_path())
+        if is_theatre:
+            self.s3_tag = "pay/wall"
+        else:
+            self.s3_tag = "wall"
+
+    def upload_to_s3(self, im: bytes):
+        s3.upload_fileobj(
+            io.BytesIO(im),
+            config.r2_bucket_name,
+            self.s3_path(),
+            ExtraArgs={"ContentType": "image/jpeg"},
+        )
+
+        # Now for the metadata
+        met = self.asset_name.replace("img", "met")
+        s3.upload_fileobj(
+            io.BytesIO(wall_metadata(self.raw_asset_name)),
+            config.r2_bucket_name,
+            self.s3_tag + "/" + met,
+            ExtraArgs={"ContentType": "text/xml"},
+        )
 
 
 class PayMovieAsset(Asset):

config-example.py

@@ -18,3 +18,12 @@
 # Sentry configuration for error logging.
 use_sentry = False
 sentry_dsn = "https://[email protected]/1"
+
+use_s3 = False
+r2_account_id = ""
+r2_bucket_name = "room-server"
+s3_connection_url = ""
+s3_access_key_id = ""
+s3_secret_access_key = ""
+
+url1_cdn_url = "http://url1.videos.wiilink24.com"

first.py

@@ -29,7 +29,7 @@ def get_config_url(service_type: str) -> str:
 def conf_first_bin_xml():
     return {
         "maint": 0,
-        "url1": get_config_url("url1"),
+        "url1": f"{config.url1_cdn_url}/",
         "url2": get_config_url("url2"),
         "url3": get_config_url("url3"),
         "eulaver": 3,

requirements.txt

@@ -3,7 +3,7 @@ lxml>=4.8
 psycopg2>=2.9
 flask_sqlalchemy>=3.0.3
 flask-login>=0.6.3
-flask-wtf>=1.0
+flask-wtf>=1.2.0
 flask-migrate>=3.1
 git+https://github.com/WiiLink24/[email protected]#egg=crc16
 pillow>=9
@@ -12,3 +12,4 @@ python-dotenv>=0.19
 sentry-sdk[flask]>=1.5.0
 SQLAlchemy-Searchable>=1.4.1
 tzdata; platform_system == 'Windows'
+boto3

room.py

@@ -1,4 +1,6 @@
+import boto3
 from flask import Flask
+from botocore.client import Config
 from flask_login import LoginManager
 from flask_migrate import Migrate
 from sentry_sdk.integrations.flask import FlaskIntegration
@@ -44,6 +46,18 @@
 
 db.configure_mappers()
 
+if config.use_s3:
+    # Finally start the S3 client
+    s3 = boto3.client(
+        "s3",
+        endpoint_url=config.s3_connection_url,
+        aws_access_key_id=config.s3_access_key_id,
+        aws_secret_access_key=config.s3_secret_access_key,
+        config=Config(signature_version="s3v4"),
+        region_name="us-east-1",
+    )
+
+
 # Import routes here.
 import first
 

theunderground/categories.py

@@ -1,4 +1,5 @@
 import os
+import config
 
 from flask import redirect, render_template, request, url_for
 from flask_login import login_required
@@ -7,7 +8,7 @@
 
 from asset_data import NormalCategoryAsset
 from models import Categories, db
-from room import app
+from room import app, s3
 from theunderground.forms import CategoryForm
 from theunderground.operations import manage_delete_item
 
@@ -105,4 +106,7 @@ def drop_category():
 @app.route("/theunderground/categories/<category>/thumbnail.jpg")
 @login_required
 def get_category_thumbnail(category):
+    if s3:
+        return redirect(f"{config.url1_cdn_url}/list/category/img/{category}.img")
+
     return NormalCategoryAsset(category).send_file()

theunderground/concierge.py

@@ -1,10 +1,19 @@
+# TODO: This is so flawed must fix later
+
+from io import BytesIO
+
 from flask import render_template, url_for, redirect
 from flask_login import login_required
 
 from models import db, ConciergeMiis, MiiMsgInfo, MiiData
 from room import app
 from theunderground.forms import ConciergeForm
 from theunderground.operations import manage_delete_item
+from room import s3
+from url1.event_today import event_today
+from url1.mii import obtain_mii, mii_met
+
+import config
 
 
 @app.route("/theunderground/concierge")
@@ -57,6 +66,8 @@ def edit_concierge(mii_id):
         db.session.add(msg7)
         db.session.add(concierge_data)
         db.session.commit()
+
+        update_mii_on_s3(mii_id)
         return redirect(url_for("list_concierge"))
 
     return render_template("concierge_edit.html", form=form)
@@ -72,3 +83,18 @@ def drop_concierge():
         return redirect(url_for("list_concierge"))
 
     return manage_delete_item(mii_id, "Concierge Mii", drop_concierge)
+
+
+def update_mii_on_s3(mii_id):
+    if s3:
+        # Update the today.xml
+        event_xml = event_today()
+        s3.upload_fileobj(BytesIO(event_xml), config.r2_bucket_name, "event/today.xml")
+
+        # Metadata
+        met = mii_met(mii_id)
+        s3.upload_fileobj(BytesIO(met), config.r2_bucket_name, f"mii/{mii_id}.met")
+
+        # Actual Mii
+        mii = obtain_mii(mii_id)
+        s3.upload_fileobj(BytesIO(mii), config.r2_bucket_name, f"mii/{mii_id}.mii")