diff --git a/backend/app/api/frames.py b/backend/app/api/frames.py index 70e0c935..2cd22aea 100644 --- 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 index 09a74311..5ffb348a 100644 --- 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 index 0df1288a..9abc9d5c 100644 --- 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 \ No newline at end of file diff --git a/backend/requirements.txt b/backend/requirements.txt index 141d7112..3c8a412d 100644 --- 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,9 +176,7 @@ 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 @@ -232,19 +184,12 @@ sqlalchemy==2.0.19 # 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 diff --git a/frontend/src/scenes/frame/panels/Assets/Asset.tsx b/frontend/src/scenes/frame/panels/Assets/Asset.tsx index 89c65f41..6fb9615e 100644 --- a/frontend/src/scenes/frame/panels/Assets/Asset.tsx +++ b/frontend/src/scenes/frame/panels/Assets/Asset.tsx @@ -1,6 +1,8 @@ import { useValues } from 'kea' import { frameLogic } from '../../frameLogic' -import { useState } from 'react' +import { useEffect, useState } from 'react' +import { apiFetch } from '../../../../utils/apiFetch' +import { Button } from '../../../../components/Button' interface AssetProps { path: string @@ -8,24 +10,51 @@ interface AssetProps { export function Asset({ path }: AssetProps) { const { frame } = useValues(frameLogic) + const isImage = path.endsWith('.png') || path.endsWith('.jpg') || path.endsWith('.jpeg') || path.endsWith('.gif') const [isLoading, setIsLoading] = useState(true) + const [asset, setAsset] = useState(null) + + useEffect(() => { + async function fetchAsset() { + setIsLoading(true) + setAsset(null) + const resource = await apiFetch(`/api/frames/${frame.id}/asset?path=${encodeURIComponent(path)}`) + const blob = await resource.blob() + setAsset(URL.createObjectURL(blob)) + setIsLoading(false) + } + fetchAsset() + }, [path]) return (
- {isImage ? ( - <> - setIsLoading(false)} - onError={() => setIsLoading(false)} - className="max-w-full" - src={`/api/frames/${frame.id}/asset?path=${encodeURIComponent(path)}`} - alt={path} - /> - {isLoading ?
Loading...
: null} - + {isLoading ? ( +
Loading...
+ ) : !asset ? ( +
Error loading asset
+ ) : isImage ? ( + setIsLoading(false)} + onError={() => setIsLoading(false)} + className="max-w-full" + src={asset} + alt={path} + /> ) : ( - <>{path} +
+
{path}
+ +
)}
) diff --git a/frontend/src/scenes/frame/panels/Assets/Assets.tsx b/frontend/src/scenes/frame/panels/Assets/Assets.tsx index 4c96c4cb..6b33eb9e 100644 --- a/frontend/src/scenes/frame/panels/Assets/Assets.tsx +++ b/frontend/src/scenes/frame/panels/Assets/Assets.tsx @@ -4,6 +4,8 @@ import { assetsLogic } from './assetsLogic' import { panelsLogic } from '../panelsLogic' import { CloudArrowDownIcon } from '@heroicons/react/24/outline' import { useState } from 'react' +import { apiFetch } from '../../../../utils/apiFetch' +import { Spinner } from '../../../../components/Spinner' function humaniseSize(size: number) { const units = ['B', 'KB', 'MB', 'GB', 'TB'] @@ -36,6 +38,7 @@ function TreeNode({ openAsset: (path: string) => void }): JSX.Element { const [expanded, setExpanded] = useState(node.path === '') + const [isDownloading, setIsDownloading] = useState(false) // If this node is a folder, display a collapsible section if (node.isFolder) { @@ -67,13 +70,27 @@ function TreeNode({ )} - {/* Download link */} { + e.preventDefault() + setIsDownloading(true) + const resource = await apiFetch(`/api/frames/${frameId}/asset?path=${encodeURIComponent(node.path)}`) + const blob = await resource.blob() + const url = URL.createObjectURL(blob) + const a = document.createElement('a') + a.href = url + a.download = node.name + a.click() + URL.revokeObjectURL(url) + setIsDownloading(false) + }} > - + {isDownloading ? ( + + ) : ( + + )} )