cabbage/ushort
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

implement new requirements

Browse Source
This commit is contained in:
cabbage
2026-03-17 18:28:47 +08:00
parent 04261c59cc
commit 714981b447
3 changed files with 245 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
urlshort.py
+183 -68
View File
@@ -3,12 +3,12 @@
URL Shortener — stdlib-only, Python 3.8+
Usage: python urlshort.py <config.json>
API:
API (all routes are prefixed with base_path, e.g. /s):
GET / Health check
POST /api/shorten Create a short URL (API key required)
GET /api/urls List all short URLs (API key required)
GET /api/urls/<code> Get info for a code (API key required)
DELETE /api/urls/<code> Delete a short URL (API key required)
GET /api/urls List all short URLs (API key required)
GET /api/urls/<code> Get info for a code (no API key required)
DELETE /api/urls/<code> Delete a short URL (API key required)
GET /<code> Redirect to original URL
"""
@@ -38,6 +38,15 @@ def load_config(path: str) -> dict:
cfg.setdefault("host", "0.0.0.0")
cfg.setdefault("port", 8080)
cfg.setdefault("db_path", "data/urlshort.db")
cfg.setdefault("retention_days", 0)
# Derive base_path from the path component of base_url.
# e.g. "http://example.com/s" → "/s"
# "http://example.com/s/" → "/s"
# "http://example.com" → ""
raw = urlparse(cfg["base_url"]).path.strip("/")
cfg["base_path"] = f"/{raw}" if raw else ""
return cfg
@@ -53,14 +62,24 @@ def init_db(db_path: str) -> None:
with sqlite3.connect(db_path) as conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS urls (
short_code TEXT PRIMARY KEY,
original_url TEXT NOT NULL,
created_at INTEGER NOT NULL,
visit_count INTEGER NOT NULL DEFAULT 0
short_code TEXT PRIMARY KEY,
original_url TEXT NOT NULL,
created_at INTEGER NOT NULL,
visit_count INTEGER NOT NULL DEFAULT 0,
retention_days INTEGER NOT NULL DEFAULT 0
)
""")
conn.commit()
# Migrate: add retention_days column if missing (existing DB).
try:
conn.execute(
"ALTER TABLE urls ADD COLUMN retention_days INTEGER NOT NULL DEFAULT 0"
)
conn.commit()
except sqlite3.OperationalError:
pass # column already exists
def db_connect(db_path: str) -> sqlite3.Connection:
conn = sqlite3.connect(db_path)
@@ -68,6 +87,18 @@ def db_connect(db_path: str) -> sqlite3.Connection:
return conn
def cleanup_expired(db_path: str) -> None:
"""Delete URLs whose retention period has elapsed."""
now = int(time.time())
with db_connect(db_path) as conn:
conn.execute(
"DELETE FROM urls WHERE retention_days > 0 "
"AND (created_at + retention_days * 86400) < ?",
(now,),
)
conn.commit()
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
@@ -97,6 +128,28 @@ class Handler(BaseHTTPRequestHandler):
# Injected by main() before the server starts.
cfg: dict = {}
# ------------------------------------------------------------------
# Routing helpers
# ------------------------------------------------------------------
def _local_path(self):
"""Return the request path with base_path prefix stripped.
Returns None (→ 404) when the request path does not start with
the configured base_path at all.
"""
raw = urlparse(self.path).path
base = self.cfg.get("base_path", "")
if base:
if raw == base or raw == base + "/":
# exact match on the prefix itself → treat as root
return "/"
if raw.startswith(base + "/"):
return raw[len(base):].rstrip("/") or "/"
# path is outside our prefix entirely
return None
return raw.rstrip("/") or "/"
# ------------------------------------------------------------------
# Logging
# ------------------------------------------------------------------
@@ -120,6 +173,19 @@ class Handler(BaseHTTPRequestHandler):
self.end_headers()
self.wfile.write(body)
def _send_plain(self, status: int, text: str) -> None:
body = text.encode("utf-8")
self.send_response(status)
self.send_header("Content-Type", "text/plain; charset=utf-8")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def _send_empty(self, status: int) -> None:
self.send_response(status)
self.send_header("Content-Length", "0")
self.end_headers()
def _error(self, status: int, message: str) -> None:
self._send_json(status, {"error": message})
@@ -130,17 +196,20 @@ class Handler(BaseHTTPRequestHandler):
self.end_headers()
# ------------------------------------------------------------------
# Auth
# Auth — purely api_key based (query param or JSON body)
# ------------------------------------------------------------------
def _authenticated(self) -> bool:
"""Accept key via Authorization: Bearer <key> header or ?api_key=."""
auth = self.headers.get("Authorization", "")
if auth.lower().startswith("bearer "):
return auth[7:].strip() == self.cfg["api_key"]
def _check_api_key(self, body: dict = None) -> bool:
"""Check api_key from query parameter or JSON body."""
qs = parse_qs(urlparse(self.path).query)
return qs.get("api_key", [""])[0] == self.cfg["api_key"]
key = qs.get("api_key", [""])[0]
if key:
return key == self.cfg["api_key"]
if body and isinstance(body, dict):
key = str(body.get("api_key", ""))
if key:
return key == self.cfg["api_key"]
return False
# ------------------------------------------------------------------
# Body
@@ -161,7 +230,12 @@ class Handler(BaseHTTPRequestHandler):
# ------------------------------------------------------------------
def do_GET(self): # noqa: N802
path = urlparse(self.path).path.rstrip("/") or "/"
cleanup_expired(self.cfg["db_path"])
path = self._local_path()
if path is None:
self._send_empty(404)
return
if path == "/":
self._send_json(200, {"status": "ok", "service": "url-shortener"})
@@ -178,37 +252,51 @@ class Handler(BaseHTTPRequestHandler):
self._handle_redirect(code)
def do_POST(self): # noqa: N802
path = urlparse(self.path).path.rstrip("/")
cleanup_expired(self.cfg["db_path"])
path = self._local_path()
if path is None:
self._send_empty(404)
return
if path == "/api/shorten":
self._handle_shorten()
else:
self._error(404, "Not found")
self._send_empty(404)
def do_DELETE(self): # noqa: N802
path = urlparse(self.path).path.rstrip("/")
cleanup_expired(self.cfg["db_path"])
path = self._local_path()
if path is None:
self._send_empty(404)
return
if path.startswith("/api/urls/"):
code = path[len("/api/urls/"):]
self._handle_delete_url(code)
else:
self._error(404, "Not found")
self._send_empty(404)
# ------------------------------------------------------------------
# Handlers
# ------------------------------------------------------------------
def _handle_shorten(self) -> None:
if not self._authenticated():
self._error(401, "Unauthorized")
return
# Parse both query parameters and JSON body
qs = parse_qs(urlparse(self.path).query)
body = self._read_json()
if body is None:
self._error(400, "Invalid JSON body")
return
original_url = str(body.get("url", "")).strip()
# Auth: check api_key from query param or body
if not self._check_api_key(body):
self._send_empty(403)
return
# url: query param takes precedence, then body
original_url = qs.get("url", [""])[0] or str(body.get("url", "")).strip()
if not original_url:
self._error(400, "Missing required field: url")
return
@@ -216,41 +304,53 @@ class Handler(BaseHTTPRequestHandler):
self._error(400, "Invalid URL — must start with http:// or https://")
return
custom_code = str(body.get("custom_code", "")).strip()
# short_length: query param, then body, then config default
raw_sl = qs.get("short_length", [None])[0]
if raw_sl is None:
raw_sl = body.get("short_length")
if raw_sl is not None:
try:
short_length = int(raw_sl)
except (ValueError, TypeError):
self._error(400, "Invalid short_length")
return
else:
short_length = self.cfg["short_length"]
# retention_days: query param, then body, then config default
raw_rd = qs.get("retention_days", [None])[0]
if raw_rd is None:
raw_rd = body.get("retention_days")
if raw_rd is not None:
try:
retention_days = int(raw_rd)
except (ValueError, TypeError):
self._error(400, "Invalid retention_days")
return
else:
retention_days = self.cfg.get("retention_days", 0)
with db_connect(self.cfg["db_path"]) as conn:
if custom_code:
row = conn.execute(
"SELECT short_code FROM urls WHERE short_code = ?", (custom_code,)
).fetchone()
if row:
self._error(409, "Custom code already in use")
return
short_code = custom_code
else:
short_code = self._unique_code(conn)
if short_code is None:
self._error(500, "Could not generate a unique short code — try again")
return
short_code = self._unique_code(conn, short_length)
if short_code is None:
self._error(500, "Could not generate a unique short code — try again")
return
created_at = int(time.time())
conn.execute(
"INSERT INTO urls (short_code, original_url, created_at) VALUES (?, ?, ?)",
(short_code, original_url, created_at),
"INSERT INTO urls (short_code, original_url, created_at, retention_days) "
"VALUES (?, ?, ?, ?)",
(short_code, original_url, created_at, retention_days),
)
conn.commit()
short_url = f"{self.cfg['base_url'].rstrip('/')}/{short_code}"
self._send_json(201, {
"short_code": short_code,
"short_url": short_url,
"original_url": original_url,
"created_at": created_at,
})
# Response: only the short URL in plain text
self._send_plain(201, short_url)
def _handle_list_urls(self) -> None:
if not self._authenticated():
self._error(401, "Unauthorized")
if not self._check_api_key():
self._send_empty(403)
return
with db_connect(self.cfg["db_path"]) as conn:
@@ -259,32 +359,41 @@ class Handler(BaseHTTPRequestHandler):
).fetchall()
base = self.cfg["base_url"].rstrip("/")
urls = [
{**dict(r), "short_url": f"{base}/{r['short_code']}"}
for r in rows
]
urls = []
for r in rows:
urls.append({
"short_code": r["short_code"],
"short_url": f"{base}/{r['short_code']}",
"original_url": r["original_url"],
"created_at": r["created_at"],
"visit_count": r["visit_count"],
})
self._send_json(200, {"count": len(urls), "urls": urls})
def _handle_get_url(self, code: str) -> None:
if not self._authenticated():
self._error(401, "Unauthorized")
return
# No API key required
with db_connect(self.cfg["db_path"]) as conn:
row = conn.execute(
"SELECT * FROM urls WHERE short_code = ?", (code,)
).fetchone()
if row is None:
self._error(404, "Short code not found")
self._send_empty(404)
return
base = self.cfg["base_url"].rstrip("/")
self._send_json(200, {**dict(row), "short_url": f"{base}/{code}"})
self._send_json(200, {
"short_code": row["short_code"],
"short_url": f"{base}/{row['short_code']}",
"original_url": row["original_url"],
"created_at": row["created_at"],
"visit_count": row["visit_count"],
})
def _handle_delete_url(self, code: str) -> None:
if not self._authenticated():
self._error(401, "Unauthorized")
# API key required — 403 if not authorized (no body)
if not self._check_api_key():
self._send_empty(403)
return
with db_connect(self.cfg["db_path"]) as conn:
@@ -292,12 +401,14 @@ class Handler(BaseHTTPRequestHandler):
"SELECT short_code FROM urls WHERE short_code = ?", (code,)
).fetchone()
if row is None:
self._error(404, "Short code not found")
# 404 if not found (no body)
self._send_empty(404)
return
conn.execute("DELETE FROM urls WHERE short_code = ?", (code,))
conn.commit()
self._send_json(200, {"message": f"Deleted '{code}'"})
# 204 on success (no body)
self._send_empty(204)
def _handle_redirect(self, code: str) -> None:
with db_connect(self.cfg["db_path"]) as conn:
@@ -319,8 +430,9 @@ class Handler(BaseHTTPRequestHandler):
# Internal utils
# ------------------------------------------------------------------
def _unique_code(self, conn: sqlite3.Connection, attempts: int = 10):
length = self.cfg["short_length"]
def _unique_code(self, conn: sqlite3.Connection, length: int = None, attempts: int = 10):
if length is None:
length = self.cfg["short_length"]
for _ in range(attempts):
code = generate_code(length)
exists = conn.execute(
@@ -342,6 +454,7 @@ def main() -> None:
cfg = load_config(sys.argv[1])
init_db(cfg["db_path"])
cleanup_expired(cfg["db_path"])
Handler.cfg = cfg
@@ -349,8 +462,10 @@ def main() -> None:
server = HTTPServer((host, port), Handler)
print(f"URL Shortener listening on http://{host}:{port}", file=sys.stderr)
print(f"Base URL : {cfg['base_url']}", file=sys.stderr)
print(f"DB path : {cfg['db_path']}", file=sys.stderr)
print(f"Base URL : {cfg['base_url']}", file=sys.stderr)
print(f"Base path : {cfg['base_path'] or '/'}", file=sys.stderr)
print(f"DB path : {cfg['db_path']}", file=sys.stderr)
print(f"Retention days : {cfg['retention_days']}", file=sys.stderr)
try:
server.serve_forever()