Merge branch 'main' into ui

# Conflicts: # .DEBIAN/control # .PKGBUILD/PKGBUILD # requirements.txt
2026-04-07 09:45:47 +03:00 · 2025-04-08 10:41:55 +02:00
parent 0d8d546545 4522425bcc
commit b6a0158ee2
16 changed files with 173 additions and 304 deletions
--- a/app/main.py
+++ b/app/main.py
@@ -23,7 +23,7 @@ from starlette.staticfiles import StaticFiles
 from starlette.templating import Jinja2Templates

 from orm import Origin, Lease, init as db_init, migrate
-from util import load_key, load_file
+from util import PrivateKey, PublicKey, load_file

 # Load variables
 load_dotenv('../version.env')
@@ -44,8 +44,8 @@ DLS_PORT = int(env('DLS_PORT', '443'))
 SITE_KEY_XID = str(env('SITE_KEY_XID', '00000000-0000-0000-0000-000000000000'))
 INSTANCE_REF = str(env('INSTANCE_REF', '10000000-0000-0000-0000-000000000001'))
 ALLOTMENT_REF = str(env('ALLOTMENT_REF', '20000000-0000-0000-0000-000000000001'))
-INSTANCE_KEY_RSA = load_key(str(env('INSTANCE_KEY_RSA', join(dirname(__file__), 'cert/instance.private.pem'))))
-INSTANCE_KEY_PUB = load_key(str(env('INSTANCE_KEY_PUB', join(dirname(__file__), 'cert/instance.public.pem'))))
+INSTANCE_KEY_RSA = PrivateKey.from_file(str(env('INSTANCE_KEY_RSA', join(dirname(__file__), 'cert/instance.private.pem'))))
+INSTANCE_KEY_PUB = PublicKey.from_file(str(env('INSTANCE_KEY_PUB', join(dirname(__file__), 'cert/instance.public.pem'))))
 TOKEN_EXPIRE_DELTA = relativedelta(days=int(env('TOKEN_EXPIRE_DAYS', 1)), hours=int(env('TOKEN_EXPIRE_HOURS', 0)))
 LEASE_EXPIRE_DELTA = relativedelta(days=int(env('LEASE_EXPIRE_DAYS', 90)), hours=int(env('LEASE_EXPIRE_HOURS', 0)))
 LEASE_RENEWAL_PERIOD = float(env('LEASE_RENEWAL_PERIOD', 0.15))
@@ -53,8 +53,8 @@ LEASE_RENEWAL_DELTA = timedelta(days=int(env('LEASE_EXPIRE_DAYS', 90)), hours=in
 CLIENT_TOKEN_EXPIRE_DELTA = relativedelta(years=12)
 CORS_ORIGINS = str(env('CORS_ORIGINS', '')).split(',') if (env('CORS_ORIGINS')) else [f'https://{DLS_URL}']

-jwt_encode_key = jwk.construct(INSTANCE_KEY_RSA.export_key().decode('utf-8'), algorithm=ALGORITHMS.RS256)
-jwt_decode_key = jwk.construct(INSTANCE_KEY_PUB.export_key().decode('utf-8'), algorithm=ALGORITHMS.RS256)
+jwt_encode_key = jwk.construct(INSTANCE_KEY_RSA.pem(), algorithm=ALGORITHMS.RS256)
+jwt_decode_key = jwk.construct(INSTANCE_KEY_PUB.pem(), algorithm=ALGORITHMS.RS256)

 # Logging
 LOG_LEVEL = logging.DEBUG if DEBUG else logging.INFO
@@ -277,10 +277,10 @@ async def _client_token():
        },
        "service_instance_public_key_configuration": {
            "service_instance_public_key_me": {
-                "mod": hex(INSTANCE_KEY_PUB.public_key().n)[2:],
-                "exp": int(INSTANCE_KEY_PUB.public_key().e),
+                "mod": hex(INSTANCE_KEY_PUB.raw().public_numbers().n)[2:],
+                "exp": int(INSTANCE_KEY_PUB.raw().public_numbers().e),
            },
-            "service_instance_public_key_pem": INSTANCE_KEY_PUB.export_key().decode('utf-8'),
+            "service_instance_public_key_pem": INSTANCE_KEY_PUB.pem().decode('utf-8'),
            "key_retention_mode": "LATEST_ONLY"
        },
    }
@@ -300,7 +300,7 @@ async def auth_v1_origin(request: Request):
    j, cur_time = json_loads((await request.body()).decode('utf-8')), datetime.now(UTC)

    origin_ref = j.get('candidate_origin_ref')
-    logging.info(f'> [  origin  ]: {origin_ref}: {j}')
+    logger.info(f'> [  origin  ]: {origin_ref}: {j}')

    data = Origin(
        origin_ref=origin_ref,
@@ -330,7 +330,7 @@ async def auth_v1_origin_update(request: Request):
    j, cur_time = json_loads((await request.body()).decode('utf-8')), datetime.now(UTC)

    origin_ref = j.get('origin_ref')
-    logging.info(f'> [  update  ]: {origin_ref}: {j}')
+    logger.info(f'> [  update  ]: {origin_ref}: {j}')

    data = Origin(
        origin_ref=origin_ref,
@@ -357,7 +357,7 @@ async def auth_v1_code(request: Request):
    j, cur_time = json_loads((await request.body()).decode('utf-8')), datetime.now(UTC)

    origin_ref = j.get('origin_ref')
-    logging.info(f'> [   code   ]: {origin_ref}: {j}')
+    logger.info(f'> [   code   ]: {origin_ref}: {j}')

    delta = relativedelta(minutes=15)
    expires = cur_time + delta
@@ -394,7 +394,7 @@ async def auth_v1_token(request: Request):
        return JSONr(status_code=400, content={'status': 400, 'title': 'invalid token', 'detail': str(e)})

    origin_ref = payload.get('origin_ref')
-    logging.info(f'> [   auth   ]: {origin_ref}: {j}')
+    logger.info(f'> [   auth   ]: {origin_ref}: {j}')

    # validate the code challenge
    challenge = b64enc(sha256(j.get('code_verifier').encode('utf-8')).digest()).rstrip(b'=').decode('utf-8')
@@ -437,7 +437,7 @@ async def leasing_v1_lessor(request: Request):

    origin_ref = token.get('origin_ref')
    scope_ref_list = j.get('scope_ref_list')
-    logging.info(f'> [  create  ]: {origin_ref}: create leases for scope_ref_list {scope_ref_list}')
+    logger.info(f'> [  create  ]: {origin_ref}: create leases for scope_ref_list {scope_ref_list}')

    lease_result_list = []
    for scope_ref in scope_ref_list:
@@ -481,7 +481,7 @@ async def leasing_v1_lessor_lease(request: Request):
    origin_ref = token.get('origin_ref')

    active_lease_list = list(map(lambda x: x.lease_ref, Lease.find_by_origin_ref(db, origin_ref)))
-    logging.info(f'> [  leases  ]: {origin_ref}: found {len(active_lease_list)} active leases')
+    logger.info(f'> [  leases  ]: {origin_ref}: found {len(active_lease_list)} active leases')

    response = {
        "active_lease_list": active_lease_list,
@@ -499,7 +499,7 @@ async def leasing_v1_lease_renew(request: Request, lease_ref: str):
    token, cur_time = __get_token(request), datetime.now(UTC)

    origin_ref = token.get('origin_ref')
-    logging.info(f'> [  renew   ]: {origin_ref}: renew {lease_ref}')
+    logger.info(f'> [  renew   ]: {origin_ref}: renew {lease_ref}')

    entity = Lease.find_by_origin_ref_and_lease_ref(db, origin_ref, lease_ref)
    if entity is None:
@@ -526,7 +526,7 @@ async def leasing_v1_lease_delete(request: Request, lease_ref: str):
    token, cur_time = __get_token(request), datetime.now(UTC)

    origin_ref = token.get('origin_ref')
-    logging.info(f'> [  return  ]: {origin_ref}: return {lease_ref}')
+    logger.info(f'> [  return  ]: {origin_ref}: return {lease_ref}')

    entity = Lease.find_by_lease_ref(db, lease_ref)
    if entity.origin_ref != origin_ref:
@@ -555,7 +555,7 @@ async def leasing_v1_lessor_lease_remove(request: Request):

    released_lease_list = list(map(lambda x: x.lease_ref, Lease.find_by_origin_ref(db, origin_ref)))
    deletions = Lease.cleanup(db, origin_ref)
-    logging.info(f'> [  remove  ]: {origin_ref}: removed {deletions} leases')
+    logger.info(f'> [  remove  ]: {origin_ref}: removed {deletions} leases')

    response = {
        "released_lease_list": released_lease_list,
@@ -577,7 +577,7 @@ async def leasing_v1_lessor_shutdown(request: Request):

    released_lease_list = list(map(lambda x: x.lease_ref, Lease.find_by_origin_ref(db, origin_ref)))
    deletions = Lease.cleanup(db, origin_ref)
-    logging.info(f'> [ shutdown ]: {origin_ref}: removed {deletions} leases')
+    logger.info(f'> [ shutdown ]: {origin_ref}: removed {deletions} leases')

    response = {
        "released_lease_list": released_lease_list,
@@ -600,7 +600,7 @@ if __name__ == '__main__':
    #
    ###

-    logging.info(f'> Starting dev-server ...')
+    logger.info(f'> Starting dev-server ...')

    ssl_keyfile = join(dirname(__file__), 'cert/webserver.key')
    ssl_certfile = join(dirname(__file__), 'cert/webserver.crt')
--- a/app/util.py
+++ b/app/util.py
@@ -1,8 +1,81 @@
 import logging

+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey, generate_private_key
+from cryptography.hazmat.primitives.serialization import load_pem_private_key, load_pem_public_key
+
 logging.basicConfig()


+class PrivateKey:
+
+    def __init__(self, data: bytes):
+        self.__key = load_pem_private_key(data, password=None)
+
+    @staticmethod
+    def from_file(filename: str) -> "PrivateKey":
+        log = logging.getLogger(__name__)
+        log.debug(f'Importing RSA-Private-Key from "{filename}"')
+
+        with open(filename, 'rb') as f:
+            data = f.read()
+
+        return PrivateKey(data=data.strip())
+
+    def raw(self) -> RSAPrivateKey:
+        return self.__key
+
+    def pem(self) -> bytes:
+        return self.__key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption()
+        )
+
+    def public_key(self) -> "PublicKey":
+        data = self.__key.public_key().public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        )
+        return PublicKey(data=data)
+
+    @staticmethod
+    def generate(public_exponent: int = 65537, key_size: int = 2048) -> "PrivateKey":
+        log = logging.getLogger(__name__)
+        log.debug(f'Generating RSA-Key')
+        key = generate_private_key(public_exponent=public_exponent, key_size=key_size)
+        data = key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption()
+        )
+        return PrivateKey(data=data)
+
+
+class PublicKey:
+
+    def __init__(self, data: bytes):
+        self.__key = load_pem_public_key(data)
+
+    @staticmethod
+    def from_file(filename: str) -> "PublicKey":
+        log = logging.getLogger(__name__)
+        log.debug(f'Importing RSA-Public-Key from "{filename}"')
+
+        with open(filename, 'rb') as f:
+            data = f.read()
+
+        return PublicKey(data=data.strip())
+
+    def raw(self) -> RSAPublicKey:
+        return self.__key
+
+    def pem(self) -> bytes:
+        return self.__key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        )
+
 def load_file(filename: str) -> bytes:
    log = logging.getLogger(f'{__name__}')
    log.debug(f'Loading contents of file "{filename}')
@@ -11,33 +84,6 @@ def load_file(filename: str) -> bytes:
    return content


-def load_key(filename: str) -> "RsaKey":
-    try:
-        # Crypto | Cryptodome on Debian
-        from Crypto.PublicKey import RSA
-        from Crypto.PublicKey.RSA import RsaKey
-    except ModuleNotFoundError:
-        from Cryptodome.PublicKey import RSA
-        from Cryptodome.PublicKey.RSA import RsaKey
-
-    log = logging.getLogger(__name__)
-    log.debug(f'Importing RSA-Key from "{filename}"')
-    return RSA.import_key(extern_key=load_file(filename), passphrase=None)
-
-
-def generate_key() -> "RsaKey":
-    try:
-        # Crypto | Cryptodome on Debian
-        from Crypto.PublicKey import RSA
-        from Crypto.PublicKey.RSA import RsaKey
-    except ModuleNotFoundError:
-        from Cryptodome.PublicKey import RSA
-        from Cryptodome.PublicKey.RSA import RsaKey
-    log = logging.getLogger(__name__)
-    log.debug(f'Generating RSA-Key')
-    return RSA.generate(bits=2048)
-
-
 class NV:
    __DRIVER_MATRIX_FILENAME = 'static/driver_matrix.json'
    __DRIVER_MATRIX: None | dict = None  # https://docs.nvidia.com/grid/ => "Driver Versions"