fix: download logic (#577)

* fix: better url parsing

* improve implementation

* fix:downloading logic

Author: baseplate-admin

src/cli/app/builder/urls.py

@@ -1,13 +1,12 @@
-from urllib.parse import urljoin, urlparse, urlunparse
+from urllib.parse import urljoin, urlparse
 from app.settings import settings
 import typer
 
 
 class UrlBuilder:
     def __init__(self, backend_url: str, frontend_url: str):
-        # Ensure schemes are present
-        self.__backend_url = self.__ensure_scheme(backend_url)
-        self.__frontend_url = self.__ensure_scheme(frontend_url)
+        self._backend_url = self.__ensure_scheme(backend_url)
+        self._frontend_url = self.__ensure_scheme(frontend_url)
 
     @staticmethod
     def __ensure_scheme(url: str) -> str:
@@ -17,65 +16,50 @@ def __ensure_scheme(url: str) -> str:
 
     @staticmethod
     def __ensure_trailing_slash(url: str) -> str:
-        parsed = urlparse(url)
-        if not parsed.path.endswith("/"):
-            new_path = parsed.path + "/"
-            parsed = parsed._replace(path=new_path)
-        return urlunparse(parsed)
+        if not url.endswith("/"):
+            return url + "/"
+        return url
 
     @classmethod
     def resolve(cls, initial_url: str | None = None) -> "UrlBuilder":
-        """
-        Logic:
-        1. If a URL is passed from CLI/Link, use it.
-        2. If not, check settings.
-        3. If still nothing, prompt user for 'Same Domain' setup.
-        """
         url_candidate = initial_url or settings.INSTANCE_URL
 
         if url_candidate:
-            # If we already have a URL, we treat it as both for simplicity
-            return cls(backend_url=url_candidate, frontend_url=url_candidate)
-
-        # Interactive Setup
-        same_domain = typer.confirm(
-            "Are the backend and frontend hosted on the same domain?", default=True
-        )
+            base = cls.__ensure_scheme(url_candidate)
+            # If "api" isn't in the URL, we assume the backend is at /api/
+            if "/api" not in base.lower():
+                base_dir = cls.__ensure_trailing_slash(base)
+                return cls(backend_url=urljoin(base_dir, "api/"), frontend_url=base_dir)
+            return cls(backend_url=base, frontend_url=base)
 
-        if same_domain:
-            domain = typer.prompt("Enter the base domain", default="chithi.dev").strip(
-                "/"
-            )
-            # Standard structure: root for frontend, /api/ for backend
-            return cls(
-                backend_url=f"https://{domain}/api/", frontend_url=f"https://{domain}/"
-            )
+        # Interactive Fallback
+        if typer.confirm("Are backend and frontend on the same domain?", default=True):
+            domain = typer.prompt("Enter base domain", default="chithi.dev").strip("/")
+            base = f"https://{domain}/"
+            return cls(backend_url=urljoin(base, "api/"), frontend_url=base)
         else:
-            frontend = typer.prompt("Enter the Frontend URL (e.g. https://chithi.dev)")
-            backend = typer.prompt(
-                "Enter the Backend URL (e.g. https://api.chithi.dev)"
-            )
+            frontend = typer.prompt("Frontend URL (e.g. https://chithi.dev)")
+            backend = typer.prompt("Backend URL (e.g. https://api.chithi.dev)")
             return cls(backend_url=backend, frontend_url=frontend)
 
     @property
-    def instance_url(self):
-        """Standardized backend base URL."""
-        return self.__ensure_trailing_slash(self.__backend_url)
+    def backend_url(self) -> str:
+        return self.__ensure_trailing_slash(self._backend_url)
 
     @property
-    def frontend_url(self):
-        """Standardized frontend base URL."""
-        return self.__ensure_trailing_slash(self.__frontend_url)
+    def frontend_url(self) -> str:
+        return self.__ensure_trailing_slash(self._frontend_url)
 
-    def upload_url(self):
-        return urljoin(self.instance_url, "upload")
+    def upload_url(self) -> str:
+        return urljoin(self.backend_url, "upload")
 
-    def config_url(self):
-        return urljoin(self.instance_url, "config")
+    def config_url(self) -> str:
+        return urljoin(self.backend_url, "config")
 
-    def download_url(self):
-        return urljoin(self.instance_url, "download/")
+    def download_url(self) -> str:
+        return urljoin(self.backend_url, "download/")
 
     def share_url(self, slug: str, key_secret: str) -> str:
-        # Share URLs should point to the Frontend UI
-        return f"{self.frontend_url}download/{slug}#{key_secret}"
+        # urljoin handles the path joining
+        base_share = urljoin(self.frontend_url, f"download/{slug}")
+        return f"{base_share}#{key_secret}"

src/cli/app/client.py

@@ -37,17 +37,19 @@ def __getattr__(self, name):
 
 
 class Client:
-    def __init__(self, instance_url: str | UrlBuilder):
-        self.urls = (
-            instance_url
-            if isinstance(instance_url, UrlBuilder)
-            else UrlBuilder(instance_url)
-        )
+    def __init__(self, urls: UrlBuilder):
+        """
+        Initialize with a UrlBuilder instance.
+        Uses backend_url for API calls and frontend_url for headers.
+        """
+        self.urls = urls
         self._session = requests.Session()
+
+        # Set headers based on the frontend URL to avoid CORS/Security issues
         self._session.headers.update(
             {
-                "Origin": self.urls.instance_url,
-                "Referer": self.urls.instance_url,
+                "Origin": self.urls.frontend_url.rstrip("/"),
+                "Referer": self.urls.frontend_url,
             }
         )
 
@@ -61,11 +63,18 @@ def close(self):
         self._session.close()
 
     @classmethod
-    def resolve(cls, instance_url: str | None = None) -> "Client":
-        """Resolve instance URL from arg/env/prompt and return Client."""
-        urls = UrlBuilder.resolve(instance_url)
+    def resolve(cls, initial_url: str | None = None) -> "Client":
+        """Resolve URLs via UrlBuilder and return a Client."""
+        urls = UrlBuilder.resolve(initial_url)
         return cls(urls)
 
+    def get_config(self) -> dict:
+        """Fetch the server configuration using the backend URL."""
+        url = self.urls.config_url()
+        response = self._session.get(url, timeout=30)
+        response.raise_for_status()
+        return response.json()
+
     def upload_file(
         self,
         file_path: Path,
@@ -74,21 +83,18 @@ def upload_file(
         expire_after: int | None = None,
     ) -> dict:
         """
-        Stream-upload *file_path* to ``POST /upload`` on the backend.
-        Returns the JSON response (contains the download key).
+        Stream-upload *file_path* to the backend.
+        Uses the resolved backend_url/upload endpoint.
         """
-        if expire_after_n_download is None:
-            expire_after_n_download = settings.EXPIRE_AFTER_N_DOWNLOAD
-
-        if expire_after is None:
-            expire_after = settings.EXPIRE_AFTER
+        # Resolve expiration settings
+        expire_n = expire_after_n_download or settings.EXPIRE_AFTER_N_DOWNLOAD
+        expire_t = expire_after or settings.EXPIRE_AFTER
 
-        if expire_after_n_download is None or expire_after is None:
+        # Fallback to server defaults if not set locally
+        if expire_n is None or expire_t is None:
             config = self.get_config()
-            if expire_after_n_download is None:
-                expire_after_n_download = config["default_number_of_downloads"]
-            if expire_after is None:
-                expire_after = config["default_expiry"]
+            expire_n = expire_n or config.get("default_number_of_downloads")
+            expire_t = expire_t or config.get("default_expiry")
 
         upload_url = self.urls.upload_url()
         display_name = filename or file_path.name
@@ -106,41 +112,52 @@ def upload_file(
         ):
             wrapped = _ProgressReader(f, pbar)
             wrapped_io = cast(BinaryIO, wrapped)
+
+            # Multipart form data
             files = {"file": (display_name, wrapped_io, "application/octet-stream")}
             data = {
                 "filename": display_name,
-                "expire_after_n_download": str(expire_after_n_download),
-                "expire_after": str(expire_after),
+                "expire_after_n_download": str(expire_n),
+                "expire_after": str(expire_t),
             }
 
             response = self._session.post(
                 upload_url,
                 data=data,
                 files=files,
-                timeout=(30, None),
+                timeout=(
+                    30,
+                    None,
+                ),  # 30s connect timeout, infinite read timeout for large files
             )
+
+            # If the backend returned HTML (redirect/error), catch it here
+            if "text/html" in response.headers.get("Content-Type", ""):
+                raise ConnectionError(
+                    f"Upload failed: Server returned HTML instead of JSON from {upload_url}"
+                )
+
             response.raise_for_status()
             return response.json()
 
-    def get_config(self) -> dict:
-        """Fetch the server configuration."""
-        url = self.urls.config_url()
-        response = self._session.get(url, timeout=30)
-        response.raise_for_status()
-        return response.json()
-
     def download_to_file(self, key: str, dest: Path) -> Path:
-        """
-        Stream-download a file and write it to *dest* with a progress bar.
-        Returns *dest*.
-        """
+        """Stream-download a file using the backend URL."""
         download_url = f"{self.urls.download_url()}{key}"
 
         with self._session.get(
             download_url, stream=True, timeout=(30, None)
         ) as response:
+            # Validation: Catch frontend HTML responses before they hit the crypto layer
+            content_type = response.headers.get("Content-Type", "")
+            if "text/html" in content_type:
+                raise ConnectionError(
+                    f"Expected binary file, but got HTML. Your API URL is likely wrong.\n"
+                    f"Attempted URL: {download_url}"
+                )
+
             response.raise_for_status()
             total = int(response.headers.get("content-length", 0)) or None
+
             with (
                 open(dest, "wb") as f,
                 tqdm(

src/cli/app/commands/download.py

@@ -2,75 +2,62 @@
 from typing import Annotated
 from urllib.parse import urlparse
 import tempfile
-import os
 import typer
-
 from app import archive, client, crypto
 from app.builder.urls import UrlBuilder
-from app.helpers.file import cleanup
 
-app = typer.Typer(help="Upload & download encrypted files via Chithi.")
+app = typer.Typer(help="Download encrypted files via Chithi.")
 
 
 @app.command()
 def download(
-    link: Annotated[str, typer.Argument(help="Chithi share URL or 'slug#key'")],
+    link: Annotated[str, typer.Argument(help="URL or 'slug#key'")],
     instance_url: Annotated[str | None, typer.Option("--url", "-u")] = None,
     password: Annotated[str | None, typer.Option("--password", "-p")] = None,
     output: Annotated[Path, typer.Option("--output", "-o")] = Path("."),
-) -> None:
-    """Download, decrypt, and extract a file."""
+):
     try:
-        slug: str
-        key_secret: str
-        inferred_url: str | None = None
+        slug, key_secret, inferred_url = "", "", None
 
-        # Case A: Full URL provided
+        # Parse the input link
         if "://" in link:
             parsed = urlparse(link)
             key_secret = parsed.fragment
-            path_parts = parsed.path.strip("/").split("/")
-
-            if len(path_parts) >= 1:
-                slug = path_parts[-1]
-            else:
-                raise ValueError("Could not extract slug from URL.")
-
-            # Reconstruct the base (e.g., https://chithi.dev)
+            path_parts = [p for p in parsed.path.split("/") if p]
+            if not key_secret or not path_parts:
+                raise ValueError(
+                    "Link must be in format: https://domain/download/SLUG#KEY"
+                )
+            slug = path_parts[-1]
             inferred_url = f"{parsed.scheme}://{parsed.netloc}"
-
-        # Case B: Just SLUG#KEY provided
         elif "#" in link:
             slug, key_secret = link.split("#", 1)
         else:
             raise ValueError("Invalid format. Use URL or SLUG#KEY")
 
-    except ValueError as e:
-        typer.echo(f"✗ Input parsing failed: {e}", err=True)
-        raise typer.Exit(1)
+        urls = UrlBuilder.resolve(initial_url=(instance_url or inferred_url))
 
-    # Resolve URLs: Priority is --url option > Link metadata > Interactive Prompt
-    urls = UrlBuilder.resolve(initial_url=(instance_url or inferred_url))
+        # Use a TemporaryDirectory for thread-safe, secure file handling
+        with tempfile.TemporaryDirectory(prefix="chithi_") as tmp_dir:
+            tmp_path = Path(tmp_dir)
+            tmp_dl = tmp_path / "encrypted.bin"
+            tmp_zip = tmp_path / "decrypted.zip"
 
-    # Process Download
-    fd, tmp_run = tempfile.mkstemp(prefix="chithi_")
-    os.close(fd)
-    tmp_dl = Path(f"{tmp_run}.dl")
-    tmp_zip = Path(f"{tmp_run}.zip")
+            # Download
+            with client.Client(urls) as c:
+                c.download_to_file(slug, tmp_dl)
 
-    try:
-        with client.Client(urls) as c:
-            c.download_to_file(slug, tmp_dl)
+            # Decrypt
+            ikm = crypto.base64url_to_ikm(key_secret)
+            crypto.decrypt(tmp_dl, tmp_zip, ikm=ikm, password=password)
 
-        ikm = crypto.base64url_to_ikm(key_secret)
-        crypto.decrypt(tmp_dl, tmp_zip, ikm=ikm, password=password)
+            #  Decompress
+            out_path = output.resolve()
+            archive.decompress(tmp_zip, out_path, password=password)
 
-        out_path = output.resolve()
-        archive.decompress(tmp_zip, out_path, password=password)
-        typer.echo(f"\n✓ Success! Extracted to {out_path}")
+            typer.secho(f"\n✓ Success! Extracted to {out_path}", fg=typer.colors.GREEN)
 
     except Exception as exc:
-        typer.echo(f"✗ Download failed: {exc}", err=True)
+        typer.secho(f"✗ Download failed: {exc}", fg=typer.colors.RED, err=True)
         raise typer.Exit(1)
-    finally:
-        cleanup(tmp_dl, tmp_zip)
+    # No 'finally' block needed! TemporaryDirectory cleans itself up automatically.

src/cli/app/commands/upload.py

@@ -10,7 +10,7 @@
 from app.helpers.file import cleanup
 from app.helpers.print import print_compact_qr
 
-app = typer.Typer(help="Upload & download encrypted files via Chithi.")
+app = typer.Typer(help="Upload encrypted files via Chithi.")
 console = Console()
 
 
@@ -74,7 +74,7 @@ def upload(
             finally:
                 cleanup(tmp_zip, tmp_enc)
 
-        # --- UI Output Logic ---
+        # UI Output Logic
         if minimal:
             # Clean output for scripts or pipes
             typer.echo(download_url)