Fix asset downloads (#114)

* get assets list * asset downloads
FrameOS · Dec 31, 2024 · 39cb19b · 39cb19b
1 parent d2b3ec7
commit 39cb19b
Show file tree

Hide file tree

Showing 6 changed files with 134 additions and 111 deletions.
diff --git a/backend/app/api/frames.py b/backend/app/api/frames.py
@@ -6,7 +6,6 @@
 from jose import JWTError, jwt
 from http import HTTPStatus
 from tempfile import NamedTemporaryFile
-from scp import SCPClient
 
 import httpx
 from fastapi import Depends, Request, HTTPException
@@ -207,20 +206,30 @@ async def api_frame_get_assets(id: int, db: Session = Depends(get_db), redis: Re
 
     assets = []
     for line in output:
-        parts = line.split(' ', 2)
-        size, mtime, path = parts
-        assets.append({
-            'path': path.strip(),
-            'size': int(size.strip()),
-            'mtime': int(mtime.strip()),
-        })
+        if line.strip():
+            parts = line.split(' ', 2)
+            size, mtime, path = parts
+            assets.append({
+                'path': path.strip(),
+                'size': int(size.strip()),
+                'mtime': int(mtime.strip()),
+            })
 
     assets.sort(key=lambda x: x['path'])
     return {"assets": assets}
 
 
 @api_with_auth.get("/frames/{id:int}/asset")
-async def api_frame_get_asset(id: int, request: Request, db: Session = Depends(get_db), redis: Redis = Depends(get_redis)):
+async def api_frame_get_asset(
+    id: int,
+    request: Request,
+    db: Session = Depends(get_db),
+    redis: Redis = Depends(get_redis)
+):
+    """
+    Download or stream an asset from the remote frame's filesystem using async SSH.
+    Uses an MD5 of the remote file to cache the content in Redis.
+    """
     frame = db.get(Frame, id)
     if frame is None:
         raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Frame not found")
@@ -234,54 +243,85 @@ async def api_frame_get_asset(id: int, request: Request, db: Session = Depends(g
         raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail="Path parameter is required")
 
     normalized_path = os.path.normpath(os.path.join(assets_path, path))
+    # Ensure the requested asset is inside the assets_path directory
     if not normalized_path.startswith(os.path.normpath(assets_path)):
         raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail="Invalid asset path")
 
     try:
         ssh = await get_ssh_connection(db, redis, frame)
         try:
+            # 1) Generate an MD5 sum of the remote file
             escaped_path = shlex.quote(normalized_path)
             command = f"md5sum {escaped_path}"
             await log(db, redis, frame.id, "stdinfo", f"> {command}")
-            stdin, stdout, stderr = ssh.exec_command(command)
-            md5sum_output = stdout.read().decode().strip()
+
+            # We'll read the MD5 from the command output
+            md5_output: list[str] = []
+            await exec_command(db, redis, frame, ssh, command, output=md5_output, log_output=False)
+            md5sum_output = "".join(md5_output).strip()
             if not md5sum_output:
                 raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Asset not found")
 
             md5sum = md5sum_output.split()[0]
             cache_key = f'asset:{md5sum}'
 
+            # 2) Check if we already have this asset cached in Redis
             cached_asset = await redis.get(cache_key)
             if cached_asset:
                 return StreamingResponse(
                     io.BytesIO(cached_asset),
                     media_type='image/png' if mode == 'image' else 'application/octet-stream',
                     headers={
-                        "Content-Disposition": f'{"attachment" if mode == "download" else "inline"}; filename={filename}'
+                        "Content-Disposition": (
+                            f'{"attachment" if mode == "download" else "inline"}; filename={filename}'
+                        )
                     }
                 )
 
-            with NamedTemporaryFile(delete=True) as temp_file:
-                with SCPClient(ssh.get_transport()) as scp:
-                    scp.get(normalized_path, temp_file.name)
-                temp_file.seek(0)
-                asset_content = temp_file.read()
-                await redis.set(cache_key, asset_content, ex=86400 * 30)
-                return StreamingResponse(
-                    io.BytesIO(asset_content),
-                    media_type='image/png' if mode == 'image' else 'application/octet-stream',
-                    headers={
-                        "Content-Disposition": f'{"attachment" if mode == "download" else "inline"}; filename={filename}'
-                    }
-                )
+            # 3) No cache found. Use asyncssh.scp to copy the remote file into a local temp file.
+            with NamedTemporaryFile(delete=False) as temp_file:
+                local_temp_path = temp_file.name
+
+            # scp from remote -> local
+            #  Note: (ssh, normalized_path) means "download from 'normalized_path' on the remote `ssh` connection"
+            import asyncssh
+            await asyncssh.scp(
+                (ssh, escaped_path),
+                local_temp_path,
+                recurse=False
+            )
+
+            # 4) Read file contents and store in Redis
+            with open(local_temp_path, "rb") as f:
+                asset_content = f.read()
+
+            await redis.set(cache_key, asset_content, ex=86400 * 30)
+
+            # Cleanup temp file
+            os.remove(local_temp_path)
+
+            # 5) Return the file to the user
+            return StreamingResponse(
+                io.BytesIO(asset_content),
+                media_type='image/png' if mode == 'image' else 'application/octet-stream',
+                headers={
+                    "Content-Disposition": (
+                        f'{"attachment" if mode == "download" else "inline"}; filename={filename}'
+                    )
+                }
+            )
+        except Exception as e:
+            print(e)
+            raise e
+
         finally:
             await remove_ssh_connection(ssh)
+
     except HTTPException:
         raise
     except Exception as e:
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=str(e))
 
-
 @api_with_auth.post("/frames/{id:int}/reset")
 async def api_frame_reset_event(id: int, redis: Redis = Depends(get_redis)):
     try:

diff --git a/backend/app/utils/ssh_utils.py b/backend/app/utils/ssh_utils.py
@@ -120,7 +120,7 @@ async def exec_command(db, redis, frame, ssh, command: str,
         # (We only store stdout in `output`, but you can also append stderr if desired.)
         if output is not None:
             stdout_data = "".join(stdout_buffer)
-            output.append(stdout_data)
+            output.extend(stdout_data.split("\n"))
 
         # Handle non-zero exit
         if exit_status != 0:

diff --git a/backend/requirements.in b/backend/requirements.in
@@ -1,16 +1,12 @@
 alembic
 arq
 asyncssh
-dacite
 email_validator
 fastapi[standard]
 honcho
 imagehash
-ipdb
 jwt
 mypy
-paramiko
-passlib
 pillow==9.5.0
 pip-tools
 pre-commit
@@ -20,11 +16,9 @@ python-jose
 redis
 requests
 ruff
-scp
 sentry-sdk[flask]
 sqlalchemy
 sqlmodel
-types-paramiko
 types-Pillow
 types-redis
 types-requests
diff --git a/backend/requirements.txt b/backend/requirements.txt
@@ -8,16 +8,8 @@ anyio==4.7.0
     #   httpx
     #   starlette
     #   watchfiles
-appnope==0.1.4
-    # via ipython
 arq==0.26.1
-asttokens==2.4.0
-    # via stack-data
 asyncssh==2.19.0
-backcall==0.2.0
-    # via ipython
-bcrypt==4.0.1
-    # via paramiko
 blinker==1.6.2
     # via
     #   flask
@@ -31,9 +23,7 @@ certifi==2023.7.22
     #   requests
     #   sentry-sdk
 cffi==1.15.1
-    # via
-    #   cryptography
-    #   pynacl
+    # via cryptography
 cfgv==3.4.0
     # via pre-commit
 charset-normalizer==3.2.0
@@ -50,15 +40,8 @@ cryptography==41.0.3
     # via
     #   asyncssh
     #   jwt
-    #   paramiko
-    #   types-paramiko
     #   types-pyopenssl
     #   types-redis
-dacite==1.8.1
-decorator==5.1.1
-    # via
-    #   ipdb
-    #   ipython
 distlib==0.3.8
     # via virtualenv
 dnspython==2.4.2
@@ -67,8 +50,6 @@ ecdsa==0.19.0
     # via python-jose
 email-validator==2.0.0.post2
     # via fastapi
-executing==1.2.0
-    # via stack-data
 fastapi==0.115.6
 fastapi-cli==0.0.6
     # via fastapi
@@ -100,13 +81,8 @@ idna==3.4
 imagehash==4.3.1
 iniconfig==2.0.0
     # via pytest
-ipdb==0.13.13
-ipython==8.15.0
-    # via ipdb
 itsdangerous==2.1.2
     # via flask
-jedi==0.19.0
-    # via ipython
 jinja2==3.1.2
     # via
     #   fastapi
@@ -122,8 +98,6 @@ markupsafe==2.1.3
     #   mako
     #   sentry-sdk
     #   werkzeug
-matplotlib-inline==0.1.6
-    # via ipython
 mdurl==0.1.2
     # via markdown-it-py
 mypy==1.13.0
@@ -140,15 +114,6 @@ packaging==23.2
     # via
     #   build
     #   pytest
-paramiko==3.3.1
-    # via scp
-parso==0.8.3
-    # via jedi
-passlib==1.7.4
-pexpect==4.8.0
-    # via ipython
-pickleshare==0.7.5
-    # via ipython
 pillow==9.5.0
     # via imagehash
 pip==24.3.1
@@ -159,12 +124,6 @@ platformdirs==4.1.0
 pluggy==1.3.0
     # via pytest
 pre-commit==3.6.0
-prompt-toolkit==3.0.39
-    # via ipython
-ptyprocess==0.7.0
-    # via pexpect
-pure-eval==0.2.2
-    # via stack-data
 pyasn1==0.6.1
     # via
     #   python-jose
@@ -178,11 +137,7 @@ pydantic==2.10.3
 pydantic-core==2.27.1
     # via pydantic
 pygments==2.16.1
-    # via
-    #   ipython
-    #   rich
-pynacl==1.5.0
-    # via paramiko
+    # via rich
 pyproject-hooks==1.0.0
     # via build
 pytest==7.4.3
@@ -213,7 +168,6 @@ rsa==4.9
 ruff==0.1.14
 scipy==1.12.0
     # via imagehash
-scp==0.14.5
 sentry-sdk==1.35.0
 setuptools==75.6.0
     # via
@@ -222,29 +176,20 @@ setuptools==75.6.0
 shellingham==1.5.4
     # via typer
 six==1.16.0
-    # via
-    #   asttokens
-    #   ecdsa
+    # via ecdsa
 sniffio==1.3.1
     # via anyio
 sqlalchemy==2.0.19
     # via
     #   alembic
     #   sqlmodel
 sqlmodel==0.0.22
-stack-data==0.6.2
-    # via ipython
 starlette==0.41.3
     # via fastapi
-traitlets==5.10.0
-    # via
-    #   ipython
-    #   matplotlib-inline
 typer==0.15.1
     # via fastapi-cli
 types-cffi==1.16.0.20240331
     # via types-pyopenssl
-types-paramiko==3.5.0.20240928
 types-pillow==10.2.0.20240822
 types-pyopenssl==24.1.0.20240722
     # via types-redis
@@ -279,8 +224,6 @@ virtualenv==20.25.0
     # via pre-commit
 watchfiles==1.0.3
     # via uvicorn
-wcwidth==0.2.6
-    # via prompt-toolkit
 websockets==14.1
     # via uvicorn
 werkzeug==2.3.6