emma-hermes/hermes_cli/managed_downloads.py

282 lines
8.8 KiB
Python

from __future__ import annotations
import base64
import hashlib
import hmac
import json
import os
import secrets
import shutil
import sqlite3
import time
from pathlib import Path
from urllib.parse import quote
import hermes_constants
import yaml
from hermes_cli.dashboard_auth import prefix as dashboard_prefix
_DEFAULT_EXPIRY_SECONDS = 7 * 24 * 60 * 60
_DOWNLOADS_PREFIX = "/downloads"
_SHORT_DOWNLOADS_PREFIX = "/downloads"
_SHORT_DOWNLOADS_TOKEN_PREFIX = "s-"
def _config_path() -> Path:
return hermes_constants.get_hermes_home() / "config.yaml"
def _load_gateway_download_public_url() -> str:
try:
cfg = yaml.safe_load(_config_path().read_text(encoding="utf-8")) or {}
except Exception:
return ""
dashboard_cfg = cfg.get("dashboard") or {}
raw = str(dashboard_cfg.get("download_public_url", "") or "").strip()
if not raw:
gateway_cfg = cfg.get("gateway") or {}
raw = str(gateway_cfg.get("download_public_url", "") or "").strip()
if not raw:
return ""
return str(dashboard_prefix._normalise_public_url(raw) or "")
def resolve_download_public_url() -> str:
env_raw = str(os.environ.get("HERMES_DOWNLOAD_PUBLIC_URL", "") or "").strip()
env_clean = str(dashboard_prefix._normalise_public_url(env_raw) or "")
if env_clean:
return env_clean
cfg_clean = _load_gateway_download_public_url()
if cfg_clean:
return cfg_clean
return str(dashboard_prefix.resolve_public_url() or "")
def _stage_dir() -> Path:
path = hermes_constants.get_hermes_home() / "state" / "public-downloads"
path.mkdir(parents=True, exist_ok=True)
return path
def stage_public_download(file_path: str) -> Path | None:
raw = str(file_path or "").strip()
if not raw:
return None
source = Path(raw).expanduser()
if not source.exists() or not source.is_file():
return None
source = source.resolve()
stage_dir = _stage_dir()
if source.is_relative_to(stage_dir):
return source
stat = source.stat()
digest = hashlib.sha256(
f"{source}|{stat.st_size}|{stat.st_mtime_ns}".encode("utf-8")
).hexdigest()[:16]
staged = stage_dir / f"{digest}-{source.name}"
if not staged.exists():
shutil.copy2(source, staged)
return staged
def _secret_path() -> Path:
return hermes_constants.get_hermes_home() / "state" / "public-download-secret.txt"
def _short_links_db_path() -> Path:
path = hermes_constants.get_hermes_home() / "state" / "public-download-links.db"
path.parent.mkdir(parents=True, exist_ok=True)
return path
def _connect_short_links_db() -> sqlite3.Connection:
conn = sqlite3.connect(_short_links_db_path())
conn.execute(
'''
CREATE TABLE IF NOT EXISTS short_links (
short_id TEXT PRIMARY KEY,
path TEXT NOT NULL,
exp INTEGER NOT NULL,
created_at INTEGER NOT NULL
)
'''
)
conn.execute("CREATE INDEX IF NOT EXISTS idx_short_links_exp ON short_links(exp)")
return conn
def _cleanup_expired_short_links(conn: sqlite3.Connection) -> None:
now = int(time.time())
conn.execute("DELETE FROM short_links WHERE exp < ?", (now,))
conn.commit()
def _persistent_download_secret() -> str:
path = _secret_path()
path.parent.mkdir(parents=True, exist_ok=True)
try:
existing = path.read_text(encoding="utf-8").strip()
except FileNotFoundError:
existing = ""
except Exception:
existing = ""
if existing:
return existing
secret = secrets.token_urlsafe(32)
path.write_text(secret, encoding="utf-8")
try:
os.chmod(path, 0o600)
except Exception:
pass
return secret
def _download_secret() -> str:
env_secret = str(os.environ.get("HERMES_DOWNLOAD_SIGNING_SECRET", "") or "").strip()
if env_secret:
return env_secret
persistent = _persistent_download_secret()
if persistent:
return persistent
env_session_secret = str(os.environ.get("HERMES_DASHBOARD_SESSION_TOKEN", "") or "").strip()
if env_session_secret:
return env_session_secret
try:
from hermes_cli import web_server
return str(getattr(web_server, "_SESSION_TOKEN", "") or "").strip()
except Exception:
return ""
def _b64url_encode(data: bytes) -> str:
return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")
def _b64url_decode(text: str) -> bytes:
padded = text + "=" * (-len(text) % 4)
return base64.urlsafe_b64decode(padded.encode("ascii"))
def create_signed_download_token(file_path: str, *, expiry_seconds: int = _DEFAULT_EXPIRY_SECONDS) -> str:
staged = stage_public_download(file_path)
secret = _download_secret()
if staged is None or not secret:
return ""
payload = {
"path": str(staged),
"exp": int(time.time()) + max(60, int(expiry_seconds or _DEFAULT_EXPIRY_SECONDS)),
}
payload_text = json.dumps(payload, separators=(",", ":"), sort_keys=True)
payload_b64 = _b64url_encode(payload_text.encode("utf-8"))
signature = hmac.new(secret.encode("utf-8"), payload_b64.encode("ascii"), hashlib.sha256).digest()
return f"{payload_b64}.{_b64url_encode(signature)}"
def create_short_download_id(file_path: str, *, expiry_seconds: int = _DEFAULT_EXPIRY_SECONDS) -> str:
staged = stage_public_download(file_path)
if staged is None:
return ""
expiry = int(time.time()) + max(60, int(expiry_seconds or _DEFAULT_EXPIRY_SECONDS))
conn = _connect_short_links_db()
try:
_cleanup_expired_short_links(conn)
for _ in range(10):
short_id = secrets.token_urlsafe(6).replace("-", "A").replace("_", "B")[:10]
if not short_id:
continue
existing = conn.execute(
"SELECT path, exp FROM short_links WHERE short_id = ?",
(short_id,),
).fetchone()
if existing is not None:
continue
conn.execute(
"INSERT INTO short_links (short_id, path, exp, created_at) VALUES (?, ?, ?, ?)",
(short_id, str(staged), expiry, int(time.time())),
)
conn.commit()
return short_id
finally:
conn.close()
return ""
def resolve_signed_download_token(token: str) -> Path:
raw = str(token or "").strip()
if not raw or "." not in raw:
if raw.startswith(_SHORT_DOWNLOADS_TOKEN_PREFIX):
return resolve_short_download_id(raw[len(_SHORT_DOWNLOADS_TOKEN_PREFIX) :])
raise ValueError("invalid_download_token")
payload_b64, signature_b64 = raw.split(".", 1)
secret = _download_secret()
if not secret:
raise ValueError("missing_download_secret")
expected = hmac.new(secret.encode("utf-8"), payload_b64.encode("ascii"), hashlib.sha256).digest()
actual = _b64url_decode(signature_b64)
if not hmac.compare_digest(expected, actual):
raise ValueError("invalid_download_signature")
payload = json.loads(_b64url_decode(payload_b64).decode("utf-8"))
exp = int(payload.get("exp", 0) or 0)
if exp < int(time.time()):
raise ValueError("expired_download_token")
target = Path(str(payload.get("path", "") or "")).expanduser().resolve()
stage_dir = _stage_dir().resolve()
if not str(target).startswith(str(stage_dir) + os.sep):
raise ValueError("download_target_outside_stage_dir")
if not target.exists() or not target.is_file():
raise FileNotFoundError(str(target))
return target
def resolve_short_download_id(short_id: str) -> Path:
raw = str(short_id or "").strip()
if not raw:
raise ValueError("invalid_short_download_id")
conn = _connect_short_links_db()
try:
row = conn.execute(
"SELECT path, exp FROM short_links WHERE short_id = ?",
(raw,),
).fetchone()
finally:
conn.close()
if row is None:
raise ValueError("invalid_short_download_id")
exp = int(row[1] or 0)
if exp < int(time.time()):
raise ValueError("expired_download_token")
target = Path(str(row[0] or "")).expanduser().resolve()
stage_dir = _stage_dir().resolve()
if not str(target).startswith(str(stage_dir) + os.sep):
raise ValueError("download_target_outside_stage_dir")
if not target.exists() or not target.is_file():
raise FileNotFoundError(str(target))
return target
def build_public_download_url(file_path: str) -> str:
base_url = str(resolve_download_public_url() or "").rstrip("/")
if not base_url:
return ""
short_id = create_short_download_id(file_path)
if not short_id:
return ""
# Preflight: prove the short link resolves before we hand it out.
resolve_short_download_id(short_id)
return f"{base_url}{_SHORT_DOWNLOADS_PREFIX}/{quote(f'{_SHORT_DOWNLOADS_TOKEN_PREFIX}{short_id}', safe='')}"