File size: 5,801 Bytes
3d5dac4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
"""
database.py — Esecuzione query SQL su database configurato.

Configurazione HF Space:
  DATABASE_URL  — URL connessione (postgresql://user:pass@host:5432/dbname oppure sqlite:///path.db)
  READ_ONLY_DB  — "true" blocca INSERT/UPDATE/DELETE/DROP (default: true per sicurezza)

Supporta:
  - PostgreSQL  (richiede psycopg2-binary in requirements.txt)
  - SQLite      (built-in Python)

Sicurezza:
  In modalità read-only (default) blocca keyword DML/DDL pericolose.
  GAP-SQL-FIX: rilevamento CTE-aware con normalizzazione whitespace + body CTE inspection.
  La vecchia blacklist regex era bypassabile con:
    - WITH del AS (DELETE FROM ...) — CTE wrapping
    - delete   from accounts — whitespace multiplo
  Fix: normalizza whitespace → collassa a singolo spazio, poi cerca nel body CTE
  anche dentro WITH ... AS (...). Zero nuove dipendenze.
  MAX 500 righe per query per evitare payload giganti.
"""
import asyncio, os, logging, re as _re
from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel

router  = APIRouter(prefix="/api/database", tags=["database"])
_logger = logging.getLogger("database")

_DANGEROUS = frozenset({"drop","truncate","delete","update","insert","alter","create","grant","revoke"})
_MAX_ROWS   = 500


class QueryRequest(BaseModel):
    sql:       str
    params:    list = []
    read_only: bool = True


def _is_dangerous(sql: str) -> str | None:
    """GAP-SQL-FIX: rilevamento keyword pericolose robusto contro CTE e multi-spazio.

    Approccio:
    1. Normalizza whitespace (collassa sequenze di spazi/newline a singolo spazio)
       → elimina bypass con spazi multipli: 'delete   from' → 'delete from'
    2. Scansiona i token separati da boundary (spazio, paren, punto-virgola)
       → rileva keyword in qualsiasi posizione, incluso dentro CTE body
    3. Per statement WITH ... AS (...), ispeziona esplicitamente il corpo delle CTE
       → rileva: WITH del AS (DELETE FROM ...), WITH x AS (DROP TABLE ...)
    """
    # Step 1: normalizza (lowercase + whitespace collassato)
    s_norm = ' '.join(sql.strip().lower().split())

    # Step 2: tokenizza su boundary SQL (spazio, (, ), ;, ,)
    tokens = _re.split(r'[\s\(\),;]+', s_norm)
    for tok in tokens:
        if tok in _DANGEROUS:
            return tok

    # Step 3: CTE body inspection esplicita (ridondante ma chiara come intento)
    # Cerca pattern: AS ( <keyword> oppure AS(<keyword>
    for m in _re.finditer(r'\bas\s*\(\s*(\w+)', s_norm):
        first_word = m.group(1).lower()
        if first_word in _DANGEROUS:
            return first_word

    return None


@router.post("/query")
async def database_query(req: QueryRequest, request: Request):
    _internal_token = os.getenv('INTERNAL_TOKEN', '')
    if _internal_token and request.headers.get('X-Internal-Token') != _internal_token:
        raise HTTPException(401, 'Unauthorized')
    db_url = os.getenv("DATABASE_URL", "").strip()
    if not db_url:
        return {
            "ok": False, "error": "DATABASE_URL non configurata.",
            "hint": "Aggiungi DATABASE_URL nelle variabili dell'HF Space (es. postgresql://user:pass@host:5432/db).",
        }

    if req.read_only or os.getenv("READ_ONLY_DB", "true").lower() != "false":
        kw = _is_dangerous(req.sql)
        if kw:
            return {
                "ok": False,
                "error": f"Query bloccata (read-only): '{kw.upper()}' non consentito.",
                "hint": "Usa read_only: false per query di scrittura, oppure imposta READ_ONLY_DB=false nell'HF Space.",
            }

    try:
        if db_url.startswith(("postgresql://", "postgres://")):
            return await _pg_query(db_url, req.sql, req.params)
        elif db_url.startswith("sqlite"):
            return await _sqlite_query(db_url, req.sql, req.params)
        else:
            return {"ok": False, "error": f"Database non supportato: {db_url[:40]}…"}
    except Exception as e:
        return {"ok": False, "error": str(e)[:400]}


async def _pg_query(db_url: str, sql: str, params: list):
    try:
        import psycopg2, psycopg2.extras
    except ImportError:
        return {"ok": False, "error": "psycopg2 non installato. Aggiungi 'psycopg2-binary' a requirements.txt."}

    def _run():
        conn = psycopg2.connect(db_url)
        try:
            cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
            cur.execute(sql, params or None)
            try:
                rows  = [dict(r) for r in cur.fetchmany(_MAX_ROWS)]
                cols  = [d.name for d in (cur.description or [])]
            except psycopg2.ProgrammingError:
                rows, cols = [], []
            conn.commit()
        finally:
            conn.close()
        return rows, cols

    rows, cols = await asyncio.to_thread(_run)
    return {
        "ok": True, "rows": rows, "columns": cols, "count": len(rows),
        "truncated": len(rows) == _MAX_ROWS,
    }


async def _sqlite_query(db_url: str, sql: str, params: list):
    import sqlite3
    path = db_url.replace("sqlite:///", "").replace("sqlite://", "") or ":memory:"

    def _run():
        conn = sqlite3.connect(path)
        conn.row_factory = sqlite3.Row
        try:
            cur  = conn.cursor()
            cur.execute(sql, params or [])
            try:
                rows = [dict(r) for r in cur.fetchmany(_MAX_ROWS)]
                cols = [d[0] for d in (cur.description or [])]
            except Exception:
                rows, cols = [], []
            conn.commit()
        finally:
            conn.close()
        return rows, cols

    rows, cols = await asyncio.to_thread(_run)
    return {
        "ok": True, "rows": rows, "columns": cols, "count": len(rows),
        "truncated": len(rows) == _MAX_ROWS,
    }