hikari_uploader/hikari/hikari.py

#!/usr/bin/env python3
import hashlib
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

import pyperclip
import requests
import typer
from requests_toolbelt import MultipartEncoder, MultipartEncoderMonitor
from rich.console import Console
from rich.progress import (BarColumn, DownloadColumn, Progress, TextColumn,
                           TimeRemainingColumn, TransferSpeedColumn)

HIKARI_BASE = "https://hikari.butaishoujo.moe"
HIKARI_URL = f"{HIKARI_BASE}/upload"

console = Console()


@dataclass
class UploadResponse:
    status: Literal["exists"] | Literal["uploaded"]
    url: str


def check_hash(filename: str, file: Path):
    h = hashlib.sha256()
    with file.open('br') as f:
        while buf := f.read(1024**2):
            h.update(buf)

    h.hexdigest()[:8]
    url = f"{HIKARI_URL}/{h.hexdigest()[:8]}"

    with file.open('br') as f:
        files = {'file': (filename, f.read(2048))}

        with requests.post(url, files=files, stream=True) as req:
            if req.status_code == 404:
                return

            req.raise_for_status()
            return UploadResponse(**req.json())


def output(copy: bool, resp: UploadResponse):
    if copy:
        pyperclip.copy(resp.url)
        action = "Copied"
    elif resp.status == "exists":
        action = "Found"
    elif resp.status == "uploaded":
        action = "Uploaded"
    else:
        action = ""

    console.print(action, resp.url, style="bold green")


@dataclass
class ProgressMonitor:
    size: int
    filename: str

    def __post_init__(self):
        self.progress = Progress(
            TextColumn("[bold blue]{task.fields[filename]}", justify="right"),
            BarColumn(bar_width=None),
            "[progress.percentage]{task.percentage:>3.1f}%",
            "•",
            DownloadColumn(),
            "•",
            TransferSpeedColumn(),
            "•",
            TimeRemainingColumn(),
        )
        self.task = self.progress.add_task("[green]Uploading",
                                           filename=self.filename,
                                           total=self.size)
        self.progress.start()

    def __call__(self, monitor: MultipartEncoderMonitor):
        progress = monitor.bytes_read
        self.progress.update(self.task, completed=progress)


def upload(file: Path,
           obstruct: bool = typer.Option(
               False, "--obstruct", "-o",
               help="Obstruct filename (by hashing)"
           ),
           filename: str = typer.Option(
               None, "--name", "-n",
               help="Rename the uploaded file."
           ),
           copy: bool = typer.Option(
               False, "--copy", "-c",
               help="Copy file URL to clipboard."
           ),
           hash: bool = typer.Option(
               True, "--hash/--no-hash", "-h/-H",
               help="Hash file before uploading to avoid having to send the whole file twice."
           )):

    if copy:
        pyperclip.copy("")

    if obstruct:
        filename = f"{hashlib.sha256(file.name.encode()).hexdigest()}{file.suffix}"
    elif filename is None:
        filename = file.name

    if hash and (resp := check_hash(filename, file)):
        return output(copy, resp)

    with file.open('br') as f:
        files = MultipartEncoder(
            fields={'file': (filename, f)}
        )
        progress_monitor = ProgressMonitor(files.len, filename)

        monitored = MultipartEncoderMonitor(files, progress_monitor)
        headers = {'Content-Type': monitored.content_type}

        with requests.post(HIKARI_URL,
                           data=monitored,  # type: ignore
                           headers=headers) as req:
            progress_monitor.progress.stop()

            req.raise_for_status()
            data = UploadResponse(**req.json())
            output(copy, data)


def main():
    typer.run(upload)