Update w private commands

app.py

@@ -1,11 +1,13 @@
 from functools import lru_cache
+import os
+import secrets
 import json
 from pathlib import Path
 import re
 from urllib.parse import quote, urlparse
 
 import duckdb
-from fastapi import FastAPI, HTTPException
+from fastapi import Depends, FastAPI, Header, HTTPException
 from huggingface_hub import hf_hub_download
 import tldextract
 
@@ -15,7 +17,7 @@ BINARY_FILENAME = "MLP-Inference/Text+GAT/mlpInfer_binaryClassifcation_dec2024_e
 REGRESSION_SOURCE = f"https://huggingface.co/datasets/{DATASET_REPO}/blob/main/{quote(REGRESSION_FILENAME, safe='/')}"
 BINARY_SOURCE = f"https://huggingface.co/datasets/{DATASET_REPO}/blob/main/{quote(BINARY_FILENAME, safe='/')}"
 
-API_VERSION = "0.2.0"
+API_VERSION = "0.2.3"
 DATA_CUTOFF_MONTH = "2024-12"
 METHOD = "GAT-TEXT"
 
@@ -35,6 +37,17 @@ _extract = tldextract.TLDExtract(include_psl_private_domains=True)
 _DOMAIN_RE = re.compile(r"^(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+[a-z0-9-]{2,}$")
 
 
+def _require_internal_access(
+    x_internal_token: str | None = Header(default=None, alias="X-Internal-Token"),
+) -> None:
+    expected = os.getenv("INTERNAL_TOKEN")
+    if not expected:
+        # Fail closed when token is not configured.
+        raise HTTPException(status_code=503, detail="Internal endpoints are disabled")
+    if not x_internal_token or not secrets.compare_digest(x_internal_token, expected):
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+
 def _canonicalize(value: str) -> str:
     if not value or not isinstance(value, str):
         raise ValueError("Domain must be a non-empty string")
@@ -187,7 +200,7 @@ def metadata():
     }
 
 
-@app.get("/label_sets")
+@app.get("/label_sets", dependencies=[Depends(_require_internal_access)], include_in_schema=False)
 def label_sets():
     return {
         "label_sets": [
@@ -225,7 +238,7 @@ def by_domain(domain: str):
     return out
 
 
-@app.get("/stats")
+@app.get("/stats", dependencies=[Depends(_require_internal_access)], include_in_schema=False)
 def stats():
     return {
         "months": [
@@ -243,7 +256,7 @@ def stats():
     }
 
 
-@app.get("/months")
+@app.get("/months", dependencies=[Depends(_require_internal_access)], include_in_schema=False)
 def months():
     return {
         "months": [