path: root/alfred/files.py

"""
Sopel plugin: 0x0 file hosting management

All commands are owner-only and work in channel or PM.

  !files                               - show usage

  !files stats                         - file count, total size, disk usage
  !files list [n]                      - last n uploads with name/size/expiry (default 5)
  !files shorten <url> [wkfo|gumx]     - shorten a URL (default: wkfo)
  !files mirror <url> [wkfo|gumx]      - mirror a remote URL to this instance
  !files prune                         - delete expired files from disk
  !files remove <filename>             - permanently remove a file
"""

import sqlite3
import subprocess
import time
from pathlib import Path

import requests as req

from sopel import plugin
from sopel.config.types import FilenameAttribute, StaticSection, ValidatedAttribute
import os as _os, sys as _sys
_alfred_dir = _os.path.dirname(_os.path.abspath(__file__))
if _alfred_dir not in _sys.path:
    _sys.path.insert(0, _alfred_dir)
import helpstrings as h


URL_ALPHABET = "DEQhd2uFteibPwq0SWBInTpA_jcZL5GKz3YCR14Ulk87Jors9vNHgfaOmMXy6Vx-"

DOMAINS = {
    "wkfo": "wk.fo",
    "gumx": "files.gumx.cc",
}


class FilesSection(StaticSection):
    url            = ValidatedAttribute("url",            default="http://127.0.0.1:5000")
    default_domain = ValidatedAttribute("default_domain", default="wkfo")
    db_path        = FilenameAttribute("db_path",         relative=False)
    data_path      = ValidatedAttribute("data_path")
    app_path       = ValidatedAttribute("app_path")


def setup(bot):
    bot.settings.define_section("files", FilesSection)


def _is_owner(bot, trigger):
    return trigger.nick == bot.settings.core.owner


def _enbase(x):
    n = len(URL_ALPHABET)
    s = ""
    while x > 0:
        s = URL_ALPHABET[int(x % n)] + s
        x = int(x // n)
    return s or URL_ALPHABET[0]


def _debase(s):
    n = len(URL_ALPHABET)
    result = 0
    for c in s:
        result = result * n + URL_ALPHABET.index(c)
    return result


def _fmt_size(n):
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024:
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} TB"


def _db(bot):
    return sqlite3.connect(bot.settings.files.db_path)


def _post(bot, data, domain=None):
    host = DOMAINS.get(domain or bot.settings.files.default_domain, "wk.fo")
    r = req.post(
        bot.settings.files.url,
        data=data,
        headers={"Host": host},
        timeout=30,
    )
    r.raise_for_status()
    return r.text.strip()


def _flask(bot, cmd, timeout=120):
    r = subprocess.run(
        ["python3", "-m", "flask", "--app", "fhost", cmd],
        capture_output=True, text=True, timeout=timeout,
        cwd=bot.settings.files.app_path,
    )
    out = r.stdout.strip()
    return out[:400] if out else "(no output)"


# ---------------------------------------------------------------------------
# !files
# ---------------------------------------------------------------------------

@plugin.command("files")
def cmd_files(bot, trigger):
    if not _is_owner(bot, trigger):
        return

    args = (trigger.group(2) or "").split()

    if not args:
        bot.say(h.FILES_TOPIC)
        return

    sub = args[0]

    if sub == "stats":
        _cmd_stats(bot)

    elif sub == "list":
        n = 5
        if len(args) > 1:
            try:
                n = max(1, min(int(args[1]), 20))
            except ValueError:
                bot.say("usage: !files list [n]")
                return
        _cmd_list(bot, n)

    elif sub == "shorten" and len(args) >= 2:
        _cmd_shorten(bot, args[1], args[2] if len(args) > 2 else None)

    elif sub == "mirror" and len(args) >= 2:
        _cmd_mirror(bot, args[1], args[2] if len(args) > 2 else None)

    elif sub == "prune":
        bot.say("pruning expired files...")
        bot.say(_flask(bot, "prune"))

    elif sub == "remove" and len(args) == 2:
        _cmd_remove(bot, args[1])

    else:
        bot.say("unknown subcommand - type !files for usage")


# ---------------------------------------------------------------------------
# stats
# ---------------------------------------------------------------------------

def _cmd_stats(bot):
    try:
        conn = _db(bot)
        live    = conn.execute(
            "SELECT COUNT(*), SUM(size) FROM file "
            "WHERE removed = 0 AND expiration IS NOT NULL"
        ).fetchone()
        permanent = conn.execute(
            "SELECT COUNT(*) FROM file WHERE expiration IS NULL AND removed = 0"
        ).fetchone()[0]
        removed = conn.execute(
            "SELECT COUNT(*) FROM file WHERE removed = 1"
        ).fetchone()[0]
        conn.close()

        count = live[0] or 0
        size  = live[1] or 0
        bot.say(f"files: {count} live | {permanent} permanent | {removed} removed | size: {_fmt_size(size)}")
    except Exception as e:
        bot.say(f"error: {e}")


# ---------------------------------------------------------------------------
# list
# ---------------------------------------------------------------------------

def _cmd_list(bot, n):
    try:
        conn = _db(bot)
        rows = conn.execute(
            "SELECT id, ext, size, expiration FROM file "
            "WHERE removed = 0 AND expiration IS NOT NULL "
            "ORDER BY id DESC LIMIT ?",
            (n,)
        ).fetchall()
        conn.close()

        if not rows:
            bot.say("files: no uploads found")
            return

        now_ms = time.time() * 1000
        for fid, ext, size, expiration in rows:
            name     = _enbase(fid) + (ext or "")
            size_str = _fmt_size(size or 0)
            hours    = max(0, (expiration - now_ms) / 3_600_000)
            bot.say(f"{name} | {size_str} | expires in {hours:.0f}h")
    except Exception as e:
        bot.say(f"error: {e}")


# ---------------------------------------------------------------------------
# shorten
# ---------------------------------------------------------------------------

def _cmd_shorten(bot, url, domain):
    if domain and domain not in DOMAINS:
        bot.say(f"unknown domain '{domain}' - use wkfo or gumx")
        return
    try:
        bot.say(_post(bot, {"shorten": url}, domain))
    except Exception as e:
        bot.say(f"error: {e}")


# ---------------------------------------------------------------------------
# mirror
# ---------------------------------------------------------------------------

def _cmd_mirror(bot, url, domain):
    if domain and domain not in DOMAINS:
        bot.say(f"unknown domain '{domain}' - use wkfo or gumx")
        return
    try:
        bot.say(f"mirroring {url}...")
        bot.say(_post(bot, {"url": url}, domain))
    except Exception as e:
        bot.say(f"error: {e}")


# ---------------------------------------------------------------------------
# remove
# ---------------------------------------------------------------------------

def _cmd_remove(bot, filename):
    filename = filename.rstrip("/").split("/")[-1]
    p    = Path(filename)
    sufs = "".join(p.suffixes[-2:])
    name = p.name[:-len(sufs) or None]

    try:
        file_id = _debase(name)
    except (ValueError, IndexError):
        bot.say(f"invalid filename: {filename}")
        return

    try:
        conn = _db(bot)
        row  = conn.execute(
            "SELECT sha256 FROM file WHERE id = ? AND removed = 0",
            (file_id,)
        ).fetchone()

        if not row:
            bot.say(f"file not found or already removed: {filename}")
            conn.close()
            return

        Path(bot.settings.files.data_path, row[0]).unlink(missing_ok=True)

        conn.execute(
            "UPDATE file SET removed = 1, expiration = NULL, mgmt_token = NULL "
            "WHERE id = ?",
            (file_id,)
        )
        conn.commit()
        conn.close()
        bot.say(f"removed {filename}")
    except Exception as e:
        bot.say(f"error: {e}")