hikari_uploader/hikari/hikari.py
2022-10-06 09:42:01 +02:00

138 lines
3.8 KiB
Python
Executable file

#!/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 = hashlib.sha256(file.name.encode()).hexdigest()
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)