add system user for devices

.gitignore

@@ -140,4 +140,6 @@ dmypy.json
 .pyre/
 
 *.toml
-!config.example.toml
+!config.example.toml
+devices/
+util_scripts/

alembic/versions/f194cc3e699f_create_system_user.py

@@ -0,0 +1,35 @@
+"""create system user
+
+Revision ID: f194cc3e699f
+Revises: 8fbc59f03986
+Create Date: 2024-08-01 12:41:28.108426
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "f194cc3e699f"
+down_revision: Union[str, None] = "8fbc59f03986"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.execute(
+        """
+        INSERT INTO `users` (`id`, `username`, `discriminator`, `avatar`, `public_flags`, `api_key`, `flags`) VALUES ('0000000000000000000', 'System', '0', NULL, '0', '0', '64');
+        """
+    )
+
+
+def downgrade() -> None:
+    op.execute(
+        """
+        DELETE FROM `users` WHERE `id` = '0000000000000000000';
+        """
+    )

getwvkeys/main.py

@@ -124,7 +124,7 @@ def wrapped_function(*args, **kwargs):
                 # check if the key is a valid user key
                 user = FlaskUser.get_user_by_api_key(db, api_key)
 
-                if not user:
+                if not user or user.flags.has(UserFlags.SYSTEM):
                     raise Forbidden("Invalid API Key")
 
                 login_user(user, remember=False)
@@ -286,8 +286,12 @@ def upload_file():
         try:
             code = gwvk.upload_device(blob, key, current_user.id)
         except Exception as e:
-            return render_template(
-                "upload.html", current_user=current_user, website_version=website_version, error=str(e)
+            logger.exception(e)
+            return (
+                render_template(
+                    "upload.html", current_user=current_user, website_version=website_version, error=str(e)
+                ),
+                400,
             )
 
         return render_template(

getwvkeys/models/Device.py

@@ -17,7 +17,7 @@
 
 import hashlib
 
-from sqlalchemy import Column, ForeignKey, Integer, String, Text, event
+from sqlalchemy import Column, ForeignKey, Integer, String, Text
 from sqlalchemy.orm import relationship
 
 from getwvkeys.models.Base import Base
@@ -50,13 +50,3 @@ def to_json(self):
             "code": self.code,
             "info": self.info,
         }
-
-
-@event.listens_for(Device, "before_insert")
-def set_info(mapper, connection, target):
-    target.code = generate_device_code(target.client_id_blob_filename, target.device_private_key)
-
-
-@event.listens_for(Device, "before_update")
-def update_info(mapper, connection, target):
-    target.code = generate_device_code(target.client_id_blob_filename, target.device_private_key)

getwvkeys/scripts.py

@@ -0,0 +1,102 @@
+import base64
+import json
+from pathlib import Path
+from typing import List
+
+import click
+from pywidevine import Device, DeviceTypes
+from sqlalchemy import create_engine, text
+from sqlalchemy.orm import Session
+
+from getwvkeys import config
+from getwvkeys.models.Device import generate_device_code
+from getwvkeys.utils import get_blob_id
+
+engine = create_engine(config.SQLALCHEMY_DATABASE_URI)
+session = Session(engine)
+
+
+@click.group()
+def cli():
+    pass
+
+
+# seed system devices, requires migration f194cc3e699f
+@cli.command()
+@click.argument("folder", type=click.Path(exists=True, file_okay=False, dir_okay=True, readable=True))
+def seed_devices(folder: click.Path):
+    cfg_out = []
+
+    system_user = session.execute(text("SELECT * FROM users WHERE id = '0000000000000000000';")).first()
+    if system_user is None:
+        raise Exception("Cannot find system user, please run the alembic migration f194cc3e699f")
+    else:
+        print("Found system user")
+
+    # check for a manifest in the folder
+    manifest_path = Path(folder) / "manifest.json"
+    if not manifest_path.exists():
+        print("ERR: Missing manifest.json")
+        exit(1)
+
+    # load the manifest
+    with open(manifest_path, "r") as f:
+        manifest: List[List[str]] = json.load(f)
+        for device in manifest:
+            client_id = device[0]
+            private_key = device[1]
+
+            client_id_blob = open(Path(folder) / client_id, "rb").read()
+            device_private_key = open(Path(folder) / private_key, "r").read()
+
+            client_id_b64 = base64.b64encode(client_id_blob).decode()
+            private_key_b64 = base64.b64encode(device_private_key.encode()).decode()
+
+            code = generate_device_code(client_id_b64, private_key_b64)
+
+            # check if device is already in the database
+            device_exists = session.execute(text("SELECT * FROM devices WHERE code = :code;"), {"code": code}).first()
+            if device_exists is not None:
+                print(f"Device {code} already exists")
+                cfg_out.append(code)
+            else:
+                info = get_blob_id(client_id_b64)
+                wvd = Device(
+                    type_=DeviceTypes.ANDROID,
+                    security_level=3,  # TODO: let user specify?
+                    flags=None,
+                    private_key=device_private_key,
+                    client_id=client_id_blob,
+                )
+
+                # insert
+                result = session.execute(
+                    text(
+                        "INSERT INTO devices (code, wvd, uploaded_by, info) VALUES (:code, :wvd, :uploaded_by, :info);"
+                    ),
+                    {
+                        "code": code,
+                        "wvd": base64.b64encode(wvd.dumps()).decode(),
+                        "uploaded_by": system_user.id,
+                        "info": info,
+                    },
+                )
+
+                if result.rowcount == 1:
+                    print(f"Device {code} created")
+                    cfg_out.append(code)
+                else:
+                    print(f"ERR: Failed to create device {code}")
+
+    session.commit()
+
+    # close
+    session.close()
+    engine.dispose()
+
+    # print json array
+    print(json.dumps(cfg_out, indent=4))
+
+
+def main():
+    cli()

getwvkeys/utils.py

@@ -52,6 +52,7 @@ class UserFlags(Enum):
     KEY_ADDING = 1 << 3
     SUSPENDED = 1 << 4
     BLACKLIST_EXEMPT = 1 << 5
+    SYSTEM = 1 << 6
 
 
 class FlagAction(Enum):

pyproject.toml

@@ -31,6 +31,7 @@ mariadb = { version = "^1.1.10", optional = true }
 mysql-connector-python = { version = "^8.4.0", optional = true }
 pywidevine = "^1.8.0"
 waitress = "^3.0.0"
+click = "^8.1.7"
 
 [tool.poetry.dev-dependencies]
 black = "^22.3.0"
@@ -39,6 +40,7 @@ isort = "^5.10.1"
 [tool.poetry.scripts]
 serve = 'getwvkeys.main:main'
 migrate = 'getwvkeys.main:run_migrations'
+gwvk = 'getwvkeys.scripts:main'
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]