app.py

from __future__ import annotations

import json
import os
import re
import sqlite3
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse

import requests
from flask import Flask, current_app, g, jsonify, redirect, render_template, request
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address

DATA_PATH = Path(__file__).parent / "data"
LEGACY_JSON_PATH = DATA_PATH / "url_store.json"

ALIAS_REGEX = re.compile(r"^[A-Za-z0-9_-]+$")
RESERVED_ALIASES = {
    "api",
    "health",
    "static",
    "favicon.svg",
    "robots.txt",
}

limiter = Limiter(get_remote_address, default_limits=["200 per day", "60 per hour"])


def _get_db() -> sqlite3.Connection:
    if "db" not in g:
        db_path = current_app.config["DATABASE_PATH"]
        DATA_PATH.mkdir(parents=True, exist_ok=True)
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        g.db = conn
    return g.db


def _close_db(exception: Exception | None = None) -> None:
    db = g.pop("db", None)
    if db is not None:
        db.close()


def init_db() -> None:
    db_path = current_app.config["DATABASE_PATH"]
    DATA_PATH.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS links (
                alias TEXT PRIMARY KEY,
                url TEXT NOT NULL,
                created_at TEXT NOT NULL,
                expires_at TEXT,
                usage_count INTEGER NOT NULL DEFAULT 0
            )
            """
        )
        conn.commit()
    finally:
        conn.close()


def _db_has_links() -> bool:
    db = _get_db()
    row = db.execute("SELECT 1 FROM links LIMIT 1").fetchone()
    return row is not None


def _migrate_json_to_db() -> None:
    if not LEGACY_JSON_PATH.exists():
        return
    if _db_has_links():
        return
    try:
        with LEGACY_JSON_PATH.open("r", encoding="utf-8") as fh:
            data = json.load(fh)
    except (json.JSONDecodeError, OSError):
        return

    if not isinstance(data, dict):
        return

    db = _get_db()
    for alias, entry in data.items():
        if not isinstance(entry, dict):
            continue
        db.execute(
            """
            INSERT OR IGNORE INTO links (alias, url, created_at, expires_at, usage_count)
            VALUES (?, ?, ?, ?, ?)
            """,
            (
                alias,
                entry.get("url"),
                entry.get("created_at"),
                entry.get("expires_at"),
                int(entry.get("usage_count") or 0),
            ),
        )
    db.commit()


def _parse_expiry(value: Optional[str]) -> Optional[str]:
    if not value:
        return None
    try:
        dt = datetime.fromisoformat(value)
    except ValueError:
        raise ValueError("Invalid expiry format. Use ISO 8601, e.g. 2025-12-31T23:59")

    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    dt = dt.astimezone(timezone.utc)
    if dt <= datetime.now(timezone.utc):
        raise ValueError("Expiry must be in the future")
    return dt.isoformat()


def _build_short_url(alias: str) -> str:
    return request.host_url.rstrip("/") + f"/{alias}"


def _create_tiny_url(target_url: str) -> Optional[str]:
    try:
        response = requests.get(
            "https://tinyurl.com/api-create.php",
            params={"url": target_url},
            timeout=8,
        )
        if response.status_code == 200:
            return response.text.strip()
    except requests.RequestException:
        return None
    return None


def _validate_url(url: str) -> Optional[str]:
    if not url:
        return "URL is required"
    if len(url) > current_app.config["MAX_URL_LENGTH"]:
        return "URL is too long"
    parsed = urlparse(url)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        return "URL must start with http:// or https://"
    return None


def _validate_alias(alias: str) -> Optional[str]:
    if not alias:
        return None
    min_len = current_app.config["MIN_ALIAS_LENGTH"]
    max_len = current_app.config["MAX_ALIAS_LENGTH"]
    if len(alias) < min_len or len(alias) > max_len:
        return f"Alias must be {min_len}-{max_len} characters"
    if alias in RESERVED_ALIASES:
        return "Alias is reserved"
    if not ALIAS_REGEX.match(alias):
        return "Alias can only contain letters, numbers, - and _"
    return None


def _alias_exists(alias: str) -> bool:
    db = _get_db()
    row = db.execute("SELECT 1 FROM links WHERE alias = ?", (alias,)).fetchone()
    return row is not None


def _insert_link(alias: str, url: str, expires_at: Optional[str]) -> None:
    db = _get_db()
    db.execute(
        """
        INSERT INTO links (alias, url, created_at, expires_at, usage_count)
        VALUES (?, ?, ?, ?, 0)
        """,
        (alias, url, datetime.now(timezone.utc).isoformat(), expires_at),
    )
    db.commit()


def _get_link(alias: str) -> Optional[sqlite3.Row]:
    db = _get_db()
    return db.execute(
        "SELECT alias, url, created_at, expires_at, usage_count FROM links WHERE alias = ?",
        (alias,),
    ).fetchone()


def _increment_usage(alias: str) -> None:
    db = _get_db()
    db.execute(
        "UPDATE links SET usage_count = usage_count + 1 WHERE alias = ?",
        (alias,),
    )
    db.commit()


def create_app(test_config: Optional[dict] = None) -> Flask:
    app = Flask(__name__)
    app.config.from_mapping(
        DATABASE_PATH=os.getenv(
            "SHARPLINK_DB_PATH", str(DATA_PATH / "url_store.db")
        ),
        MAX_URL_LENGTH=2048,
        MIN_ALIAS_LENGTH=3,
        MAX_ALIAS_LENGTH=32,
        MAX_CONTENT_LENGTH=16 * 1024,
        RATELIMIT_STORAGE_URI=os.getenv("RATELIMIT_STORAGE_URI", "memory://"),
        RATELIMIT_ENABLED=os.getenv("RATELIMIT_ENABLED", "true").lower()
        not in {"0", "false", "no"},
    )

    if test_config:
        app.config.update(test_config)

    if app.config.get("TESTING"):
        app.config.setdefault("RATELIMIT_ENABLED", False)
        app.config.setdefault("RATELIMIT_STORAGE_URI", "memory://")

    limiter.init_app(app)
    app.teardown_appcontext(_close_db)

    with app.app_context():
        init_db()
        _migrate_json_to_db()

    @app.route("/", methods=["GET"])
    def index():
        return render_template("index.html")

    @app.route("/api/shorten", methods=["POST"])
    @limiter.limit("10 per minute")
    def shorten():
        if not request.is_json:
            return jsonify({"error": "Request must be JSON"}), 400

        payload = request.get_json(silent=True) or {}
        original_url = (payload.get("url") or "").strip()
        custom_alias = (payload.get("alias") or "").strip()
        expires_at_raw = (payload.get("expires_at") or "").strip()
        use_tiny = bool(payload.get("use_tiny"))

        url_error = _validate_url(original_url)
        if url_error:
            return jsonify({"error": url_error}), 400

        alias_error = _validate_alias(custom_alias)
        if alias_error:
            return jsonify({"error": alias_error}), 400

        alias = custom_alias or _generate_alias()
        if _alias_exists(alias):
            return jsonify({"error": "Alias already exists, choose another."}), 409

        try:
            expires_at = _parse_expiry(expires_at_raw)
        except ValueError as exc:
            return jsonify({"error": str(exc)}), 400

        try:
            _insert_link(alias, original_url, expires_at)
        except sqlite3.IntegrityError:
            return jsonify({"error": "Alias already exists, choose another."}), 409

        short_url = _build_short_url(alias)
        tiny_url = _create_tiny_url(original_url) if use_tiny else None

        return jsonify(
            {
                "alias": alias,
                "short_url": short_url,
                "expires_at": expires_at,
                "tiny_url": tiny_url,
            }
        )

    @app.route("/<alias>", methods=["GET"])
    def redirect_alias(alias: str):
        entry = _get_link(alias)
        if not entry:
            return render_template(
                "not_found.html", message="Short link not found."  # pragma: no cover
            ), 404

        expires_at = entry["expires_at"]
        if expires_at:
            expiry_dt = datetime.fromisoformat(expires_at)
            if expiry_dt <= datetime.now(timezone.utc):
                return render_template(
                    "expired.html", message="This link has expired."  # pragma: no cover
                ), 410

        _increment_usage(alias)
        return redirect(entry["url"], code=302)

    @app.route("/health", methods=["GET"])
    @limiter.exempt
    def health():
        return jsonify({"status": "ok"})

    return app


def _generate_alias(length: int = 6) -> str:
    while True:
        alias = uuid.uuid4().hex[:length]
        if not _alias_exists(alias):
            return alias


app = create_app()


if __name__ == "__main__":
    port = int(os.getenv("PORT", "5000"))
    app.run(host="0.0.0.0", port=port, debug=False)