ghostbase-cli 0.1.1__tar.gz → 0.1.3__tar.gz

{ghostbase_cli-0.1.1 → ghostbase_cli-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ghostbase-cli
-Version: 0.1.1
+Version: 0.1.3
 Summary: A fast, keyboard-driven terminal database studio for SQLite, PostgreSQL & MySQL
 Author-email: 4yk5 <contact.4yk5@proton.me>
 License-Expression: MIT

{ghostbase_cli-0.1.1 → ghostbase_cli-0.1.3}/app.py

@@ -342,7 +342,7 @@ class GhostBaseApp(App):
             yield Label(f"[{self.db.db_type_label()}]", id="db-type-label")
             yield Label(db_label, id="db-name-label")
             yield Label(safe_text, id="safe-badge", classes=safe_class)
-            yield Label("v0.1.1", id="version-label")
+            yield Label("v0.1.3", id="version-label")
 
         with Horizontal(id="main-container"):
             with Vertical(id="sidebar"):
@@ -389,14 +389,22 @@ class GhostBaseApp(App):
     # Lifecycle
 
     async def on_mount(self):
+        # Configure database engine level Read-Only mode if safe_mode is enabled (fail-closed)
+        if self.safe_mode:
+            try:
+                self.db.set_read_only(True)
+            except Exception as e:
+                self.safe_mode = False
+                self.notify(f"Failed to enable Read-Only mode on database: {e}", severity="error")
+
         tables = self.db.get_tables()
-        # Only populate demo tables when explicitly requested, in-memory, or default sample.db
+        # Only populate demo tables when explicitly requested, in-memory, or default sample.db AND NOT in safe_mode
         is_default_sample = (
             self.populate_sample
             or self.db_path == ":memory:"
             or os.path.basename(os.path.abspath(self.db_path)) == "sample.db"
         )
-        if not tables and self.db.db_type == "sqlite" and is_default_sample:
+        if not tables and self.db.db_type == "sqlite" and is_default_sample and not self.safe_mode:
             self._create_sample_data()
             tables = self.db.get_tables()
 
@@ -408,8 +416,9 @@ class GhostBaseApp(App):
         if self.all_tables:
             first = self.all_tables[0]
             self.selected_table = first
+            quoted = self.db.quote_table_name(first)
             self.query_one("#query-box", TextArea).text = (
-                f"SELECT * FROM {first} LIMIT 50;"
+                f"SELECT * FROM {quoted} LIMIT 50;"
             )
             self.action_run_query()
             self._show_schema(first)
@@ -586,8 +595,9 @@ class GhostBaseApp(App):
         self._pending_destructive_query = None
         self._pending_destructive_time = 0.0
 
+        quoted = self.db.quote_table_name(t)
         self.query_one("#query-box", TextArea).text = (
-            f"SELECT * FROM {t} LIMIT 50;"
+            f"SELECT * FROM {quoted} LIMIT 50;"
         )
         self.action_run_query()
         self._show_schema(t, is_view=is_view)
@@ -673,9 +683,10 @@ class GhostBaseApp(App):
             self.last_columns = res["columns"]
             self.last_rows    = res["rows"]
 
-            # Detect target table for export if simple SELECT FROM table
-            m = re.search(r"\bFROM\s+[`\"\[]?([a-zA-Z0-9_\.]+)[`\"\]]?", query, flags=re.IGNORECASE)
-            self.last_query_table = m.group(1).split(".")[-1] if m else (self.selected_table or "query_result")
+            # Detect target table for export from final SELECT statement (ignoring CTE definitions)
+            main_q = re.sub(r"^WITH\s+[\s\S]*?\)\s*", "", query, flags=re.IGNORECASE)
+            m = re.search(r"\bFROM\s+([`\"\[]?[a-zA-Z0-9_\.]+(?:\"|\'|\]|`)?)", main_q, flags=re.IGNORECASE)
+            self.last_query_table = m.group(1).strip('`"[]') if m else (self.selected_table or "query_result")
 
             dt.add_columns(*self.last_columns)
             for row in self.last_rows:
@@ -702,9 +713,46 @@ class GhostBaseApp(App):
             else:
                 res_cnt.update(f"({rc} rows)")
 
+            self._update_status_bar(row_count=rc, ms=ms)
+        elif res.get("is_returning"):
+            # DML with RETURNING: Render returned rows in DataTable, but treat as write operation
+            self.last_columns = res["columns"]
+            self.last_rows    = res["rows"]
+
+            # Extract target table from INSERT/UPDATE
+            m = re.search(r"\b(?:INTO|UPDATE)\s+([`\"\[]?[a-zA-Z0-9_\.]+(?:\"|\'|\]|`)?)", query, flags=re.IGNORECASE)
+            self.last_query_table = m.group(1).strip('`"[]') if m else (self.selected_table or "query_result")
+
+            dt.add_columns(*self.last_columns)
+            for row in self.last_rows:
+                styled = []
+                for val in row:
+                    if val is None:
+                        styled.append(Text("NULL", style="dim #6e7681"))
+                    elif isinstance(val, (int, float)):
+                        styled.append(Text(str(val), style="#79c0ff"))
+                    else:
+                        s = str(val).replace("\r\n", " ").replace("\n", " ").replace("\t", " ")
+                        if len(s) > 200:
+                            s = s[:197] + "..."
+                        styled.append(Text(s, style="#c9d1d9"))
+                dt.add_row(*styled)
+
+            aff = res.get("affected_rows", len(self.last_rows))
+            rc = res.get("row_count", len(self.last_rows))
+            ms = res["elapsed_ms"]
+            status.update(f"[bold green]OK[/] ({ms}ms)")
+            res_cnt.update(f"({rc} rows returned, {aff} affected)")
+            self.notify(f"Done — {aff} row(s) affected (RETURNING {rc} rows)")
+
+            # Refresh schema & row counts
+            self.all_tables = self.db.get_tables()
+            self.all_views  = self.db.get_views()
+            self._refresh_row_counts()
+            self._refresh_sidebar(self.all_tables, self.all_views)
             self._update_status_bar(row_count=rc, ms=ms)
         else:
-            # Clear export cache so previous SELECT results are not mistakenly exported
+            # Pure DML / DDL statement -> clear export cache so previous SELECT results are not mistakenly exported
             self.last_columns = []
             self.last_rows    = []
             self.last_query_table = None
@@ -766,11 +814,19 @@ class GhostBaseApp(App):
         self.notify("Execution Plan analyzed successfully")
 
     def action_toggle_safe_mode(self):
-        """Toggle Production Safe Mode (Ctrl+S)"""
-        self.safe_mode = not self.safe_mode
+        """Toggle Production Safe Mode (Ctrl+S) - Fail-closed"""
+        new_mode = not self.safe_mode
         self._pending_destructive_query = None
         self._pending_destructive_time = 0.0
 
+        # Enforce/revoke database session level Read-Only (Fail-closed)
+        try:
+            self.db.set_read_only(new_mode)
+            self.safe_mode = new_mode
+        except Exception as e:
+            self.notify(f"Cannot enable Safe Mode: {e}", severity="error", timeout=5)
+            return
+
         badge = self.query_one("#safe-badge", Label)
         if self.safe_mode:
             badge.update("[SAFE: ON]")
@@ -858,7 +914,7 @@ class GhostBaseApp(App):
             else os.getcwd()
         )
         out = os.path.join(db_dir, "ghostbase_export.sql")
-        Exporter.to_sql_inserts(table_name, self.last_columns, self.last_rows, out)
+        Exporter.to_sql_inserts(table_name, self.last_columns, self.last_rows, out, dialect=self.db.db_type)
         self.notify(f"Saved: {out}", timeout=5)
 
     # Status bar

{ghostbase_cli-0.1.1 → ghostbase_cli-0.1.3}/db_engine.py

@@ -65,16 +65,58 @@ class DatabaseEngine:
                 "Run:  pip install pymysql"
             )
         parsed = urllib.parse.urlparse(url)
+        # Use SSDictCursor (Server-Side streaming cursor) to protect client memory on massive queries
         self.conn = pymysql.connect(
             host=parsed.hostname or "localhost",
             port=parsed.port or 3306,
             user=parsed.username or "root",
             password=parsed.password or "",
             database=parsed.path.lstrip("/"),
-            cursorclass=pymysql.cursors.DictCursor,
+            cursorclass=pymysql.cursors.SSDictCursor,
             autocommit=True,
         )
 
+    def set_read_only(self, enabled: bool) -> bool:
+        """
+        Configure database engine session-level Read-Only mode.
+        Returns True if successful. Raises RuntimeError if enabling fails (fail-closed).
+        """
+        if not self.conn:
+            return False
+        try:
+            cur = self.conn.cursor()
+            if self.db_type == "sqlite":
+                val = "ON" if enabled else "OFF"
+                cur.execute(f"PRAGMA query_only = {val};")
+            elif self.db_type == "postgresql":
+                mode = "ON" if enabled else "OFF"
+                cur.execute(f"SET default_transaction_read_only = {mode};")
+            elif self.db_type == "mysql":
+                mode = "READ ONLY" if enabled else "READ WRITE"
+                cur.execute(f"SET SESSION TRANSACTION {mode};")
+            return True
+        except Exception as e:
+            if enabled:
+                raise RuntimeError(f"Database engine failed to activate Read-Only mode: {e}")
+            return False
+
+    def quote_table_name(self, table_name: str) -> str:
+        """Properly quote table identifier for queries based on database engine."""
+        if not table_name:
+            return ""
+        if self.db_type == "postgresql":
+            if "." in table_name:
+                s, t = table_name.split(".", 1)
+                return f'"{s}"."{t}"'
+            return f'"{table_name}"'
+        elif self.db_type == "mysql":
+            if "." in table_name:
+                s, t = table_name.split(".", 1)
+                return f"`{s}`.`{t}`"
+            return f"`{table_name}`"
+        else:  # sqlite
+            return f"[{table_name}]"
+
     # --- Schema discovery ---
 
     def get_tables(self) -> List[str]:
@@ -266,16 +308,28 @@ class DatabaseEngine:
             if self.db_type == "sqlite":
                 cur.execute(f"SELECT COUNT(*) FROM [{table_name}];")
             elif self.db_type == "postgresql":
-                clean_name = table_name.split(".", 1)[-1]
+                schema = "public"
+                tbl = table_name
+                if "." in table_name:
+                    schema, tbl = table_name.split(".", 1)
+
                 if fast_estimate:
                     try:
-                        cur.execute("SELECT reltuples::bigint FROM pg_class WHERE relname = %s;", (clean_name,))
+                        cur.execute(
+                            """
+                            SELECT c.reltuples::bigint
+                            FROM pg_class c
+                            JOIN pg_namespace n ON n.oid = c.relnamespace
+                            WHERE n.nspname = %s AND c.relname = %s;
+                            """,
+                            (schema, tbl),
+                        )
                         r = cur.fetchone()
                         if r and r[0] is not None and r[0] >= 0:
                             return int(r[0])
                     except Exception:
                         pass
-                cur.execute(f'SELECT COUNT(*) FROM "{table_name}";')
+                cur.execute(f'SELECT COUNT(*) FROM "{schema}"."{tbl}";')
             else:
                 if fast_estimate:
                     try:
@@ -299,20 +353,47 @@ class DatabaseEngine:
 
     # --- Query safety and formatting ---
 
+    @staticmethod
+    def _mask_strings_and_literals(sql: str) -> Tuple[str, List[str]]:
+        """
+        Extract and mask all string literals and quoted identifiers:
+        - PostgreSQL dollar-quoted strings with exact tag matching: $tag$ ... $tag$ (using backreference \2)
+        - Standard single-quoted strings: '...' (with '' escaping)
+        - Standard double-quoted strings / identifiers: "..."
+        - MySQL backtick identifiers: `...`
+        """
+        literals = []
+        def _repl(m):
+            idx = len(literals)
+            literals.append(m.group(0))
+            return f" ___GB_LIT_{idx}___ "
+
+        # Matches:
+        # 1. Plain untagged dollar-quotes: $$...$$
+        # 2. Tagged dollar-quotes with strict matching backreference: $tag$...$tag$
+        # 3. Standard single-quoted strings: '...'
+        # 4. Standard double-quoted strings/identifiers: "..."
+        # 5. MySQL backticks: `...`
+        pattern = r"(\$\$[\s\S]*?\$\$|\$([a-zA-Z_]\w*)\$[\s\S]*?\$\2\$|'(?:''|[^'])*'|\"(?:\"\"|[^\"])*\"|`(?:``|[^`])*`)"
+        masked = re.sub(pattern, _repl, sql)
+        return masked, literals
+
     def check_safety(self, query: str, is_safe_mode: bool = False) -> Tuple[bool, str]:
         """
         Check query safety.
-        In Safe Mode: Enforces strict Read-Only operations.
-        In Normal Mode: Intercepts destructive queries lacking constraints (e.g. DELETE/UPDATE without WHERE).
-        Comments (-- and /* */) and string literals are stripped to prevent syntax bypasses.
+        - Strings and dollar-quotes are masked FIRST so inner text ('WHERE ...') does not bypass checks.
+        - Comments (-- and /* */) are stripped second.
+        - In Safe Mode: Enforces strict Read-Only operations, blocking all DML, DDL, DO blocks, CALL, and SET configuration changes.
+        - In Normal Mode: Intercepts unconstrained destructive queries.
         """
-        # 1. Strip block comments /* ... */ and line comments -- ...
-        q_no_comments = re.sub(r"/\*[\s\S]*?\*/", " ", query)
+        # 1. Mask all string literals and dollar-quoted strings FIRST
+        masked, _ = self._mask_strings_and_literals(query)
+
+        # 2. Strip block comments /* ... */ and line comments -- ...
+        q_no_comments = re.sub(r"/\*[\s\S]*?\*/", " ", masked)
         q_no_comments = re.sub(r"--[^\r\n]*", " ", q_no_comments)
 
-        # 2. Strip string literals & identifiers so keywords inside literals don't cause false detections
-        q_clean = re.sub(r"('(?:''|[^'])*'|\"(?:\"\"|[^\"])*\"|`(?:``|[^`])*`)", " ", q_no_comments)
-        q_clean = q_clean.strip()
+        q_clean = q_no_comments.strip()
         if not q_clean:
             return False, ""
 
@@ -332,10 +413,18 @@ class DatabaseEngine:
                 (r"\bGRANT\b", "GRANT statement"),
                 (r"\bREVOKE\b", "REVOKE statement"),
                 (r"\bVACUUM\b", "VACUUM statement"),
+                (r"\bNEXTVAL\s*\(", "NEXTVAL sequence advance"),
+                (r"\bINTO\s+OUTFILE\b", "INTO OUTFILE export"),
+                (r"\bINTO\s+DUMPFILE\b", "INTO DUMPFILE export"),
+                (r"\bSET\b", "SET configuration change"),
+                (r"\bRESET\b", "RESET configuration change"),
+                (r"\bDO\b", "DO procedural block"),
+                (r"\bCALL\b", "CALL stored procedure"),
+                (r"\bMERGE\b", "MERGE statement"),
             ]
             for pattern, desc in write_keywords:
                 if re.search(pattern, q_upper):
-                    return True, f"Safe Mode is active: write operations ({desc}) are strictly blocked"
+                    return True, f"Safe Mode is active: write or state-altering operations ({desc}) are strictly blocked"
 
             if re.search(r"\bPRAGMA\b\s+\w+\s*=", q_upper):
                 return True, "Safe Mode is active: PRAGMA configuration writes are blocked"
@@ -357,8 +446,12 @@ class DatabaseEngine:
         return False, ""
 
     def explain_query(self, query: str) -> Dict[str, Any]:
-        clean_q = re.sub(r"/\*[\s\S]*?\*/", " ", query)
+        masked, literals = self._mask_strings_and_literals(query)
+        clean_q = re.sub(r"/\*[\s\S]*?\*/", " ", masked)
         clean_q = re.sub(r"--[^\r\n]*", " ", clean_q).strip().rstrip(";")
+        for idx, orig in enumerate(literals):
+            clean_q = clean_q.replace(f"___GB_LIT_{idx}___", orig)
+
         if not clean_q:
             return {"success": False, "error": "Query is empty"}
 
@@ -374,15 +467,7 @@ class DatabaseEngine:
         if not query or not query.strip():
             return query
 
-        # Mask string literals ('...'), double-quoted strings ("..."), and backticks (`...`)
-        literals = []
-        def _mask_literal(m):
-            idx = len(literals)
-            literals.append(m.group(0))
-            return f"___GB_LIT_{idx}___"
-
-        lit_pattern = r"('(?:''|[^'])*'|\"(?:\"\"|[^\"])*\"|`(?:``|[^`])*`)"
-        masked = re.sub(lit_pattern, _mask_literal, query)
+        masked, literals = DatabaseEngine._mask_strings_and_literals(query)
 
         lines = masked.splitlines()
         comment_lines = []
@@ -440,8 +525,27 @@ class DatabaseEngine:
 
     def execute_query(self, query: str, max_rows: int = 1000) -> Dict[str, Any]:
         start = time.perf_counter()
+        cur = None
         try:
-            cur = self.conn.cursor()
+            # Check if query is DML/DDL write operation (even if it has RETURNING)
+            masked, _ = self._mask_strings_and_literals(query)
+            clean_q = re.sub(r"/\*[\s\S]*?\*/", " ", masked)
+            clean_q = re.sub(r"--[^\r\n]*", " ", clean_q).strip().upper()
+            is_write_dml = bool(
+                re.search(r"\b(INSERT\s+INTO|UPDATE\s+|DELETE\s+FROM|DELETE\s+|REPLACE\s+INTO|CREATE\s+|DROP\s+|ALTER\s+|TRUNCATE\s+|DO\s+|CALL\s+|MERGE\s+)\b", clean_q)
+            )
+
+            # Use server-side streaming cursor for PostgreSQL read queries to prevent client RAM overflow
+            if self.db_type == "postgresql" and not is_write_dml:
+                stream_name = f"gb_stream_{int(time.perf_counter() * 1000000)}"
+                try:
+                    cur = self.conn.cursor(name=stream_name)
+                    cur.itersize = max_rows + 1
+                except Exception:
+                    cur = self.conn.cursor()
+            else:
+                cur = self.conn.cursor()
+
             cur.execute(query)
             elapsed_ms = (time.perf_counter() - start) * 1000
 
@@ -469,9 +573,11 @@ class DatabaseEngine:
                     "columns": columns,
                     "rows": rows,
                     "row_count": len(rows),
+                    "affected_rows": max(cur.rowcount, len(rows)) if is_write_dml else 0,
                     "truncated": truncated,
                     "elapsed_ms": round(elapsed_ms, 2),
-                    "is_select": True,
+                    "is_select": not is_write_dml,
+                    "is_returning": is_write_dml,
                 }
             else:
                 if self.db_type in ("sqlite", "postgresql"):
@@ -481,6 +587,7 @@ class DatabaseEngine:
                     "affected_rows": max(cur.rowcount, 0),
                     "elapsed_ms": round(elapsed_ms, 2),
                     "is_select": False,
+                    "is_returning": False,
                 }
         except Exception as e:
             elapsed_ms = (time.perf_counter() - start) * 1000
@@ -494,6 +601,12 @@ class DatabaseEngine:
                 "error": str(e),
                 "elapsed_ms": round(elapsed_ms, 2),
             }
+        finally:
+            if cur:
+                try:
+                    cur.close()
+                except Exception:
+                    pass
 
     # --- Helpers ---
 

{ghostbase_cli-0.1.1 → ghostbase_cli-0.1.3}/exporter.py

@@ -1,9 +1,30 @@
 import csv
 import json
+import decimal
+import datetime
+import uuid
 from typing import List, Any
 
 
 class Exporter:
+    @staticmethod
+    def _json_serial(obj: Any) -> Any:
+        """Serialize non-standard types (Decimal, datetime, UUID, bytes) for JSON export."""
+        if isinstance(obj, (datetime.date, datetime.datetime, datetime.time)):
+            return obj.isoformat()
+        if isinstance(obj, decimal.Decimal):
+            return int(obj) if (obj % 1 == 0) else float(obj)
+        if isinstance(obj, uuid.UUID):
+            return str(obj)
+        if isinstance(obj, (bytes, bytearray)):
+            try:
+                return obj.decode("utf-8")
+            except Exception:
+                return obj.hex()
+        if hasattr(obj, "__str__"):
+            return str(obj)
+        raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
+
     @staticmethod
     def _dedup_columns(columns: List[str]) -> List[str]:
         seen = {}
@@ -19,8 +40,15 @@ class Exporter:
         return deduped
 
     @staticmethod
-    def _escape_ident(ident: str) -> str:
+    def _escape_ident(ident: str, dialect: str = "sqlite") -> str:
         s = str(ident).strip()
+        q = "`" if str(dialect).lower() == "mysql" else '"'
+        escaped_q = q * 2
+
+        if "." in s:
+            parts = s.split(".")
+            return ".".join(f"{q}{p.replace(q, escaped_q)}{q}" for p in parts)
+
         reserved = {
             "ALL", "ALTER", "AND", "ANY", "AS", "ASC", "BETWEEN", "BY", "CASE",
             "CHECK", "COLUMN", "CONSTRAINT", "CREATE", "CROSS", "CURRENT_DATE",
@@ -34,8 +62,8 @@ class Exporter:
             "USER", "VALUES", "VIEW", "WHEN", "WHERE", "WITH"
         }
         if not s.isidentifier() or s.upper() in reserved:
-            escaped = s.replace('"', '""')
-            return f'"{escaped}"'
+            escaped = s.replace(q, escaped_q)
+            return f"{q}{escaped}{q}"
         return s
 
     @staticmethod
@@ -52,7 +80,7 @@ class Exporter:
         for r in rows:
             data.append(dict(zip(clean_cols, r)))
         with open(output_path, "w", encoding="utf-8") as f:
-            json.dump(data, f, indent=2, ensure_ascii=False)
+            json.dump(data, f, indent=2, ensure_ascii=False, default=Exporter._json_serial)
 
     @staticmethod
     def to_markdown(columns: List[str], rows: List[List[Any]], output_path: str):
@@ -70,13 +98,13 @@ class Exporter:
                 f.write("| " + " | ".join(clean_cells) + " |\n")
 
     @staticmethod
-    def to_sql_inserts(table_name: str, columns: List[str], rows: List[List[Any]], output_path: str):
+    def to_sql_inserts(table_name: str, columns: List[str], rows: List[List[Any]], output_path: str, dialect: str = "sqlite"):
         table = (table_name or "exported_table").strip()
-        escaped_table = Exporter._escape_ident(table)
-        escaped_cols = [Exporter._escape_ident(c) for c in columns]
+        escaped_table = Exporter._escape_ident(table, dialect=dialect)
+        escaped_cols = [Exporter._escape_ident(c, dialect=dialect) for c in columns]
         cols_str = ", ".join(escaped_cols)
         with open(output_path, "w", encoding="utf-8") as f:
-            f.write(f"-- GhostBase SQL Export: {table}\n")
+            f.write(f"-- GhostBase SQL Export: {table} ({dialect})\n")
             f.write(f"-- Total: {len(rows)} row(s)\n\n")
             for r in rows:
                 vals = []
@@ -85,7 +113,7 @@ class Exporter:
                         vals.append("NULL")
                     elif isinstance(v, bool):
                         vals.append("TRUE" if v else "FALSE")
-                    elif isinstance(v, (int, float)):
+                    elif isinstance(v, (int, float, decimal.Decimal)):
                         vals.append(str(v))
                     else:
                         escaped = str(v).replace("'", "''")

{ghostbase_cli-0.1.1 → ghostbase_cli-0.1.3}/ghostbase_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ghostbase-cli
-Version: 0.1.1
+Version: 0.1.3
 Summary: A fast, keyboard-driven terminal database studio for SQLite, PostgreSQL & MySQL
 Author-email: 4yk5 <contact.4yk5@proton.me>
 License-Expression: MIT

{ghostbase_cli-0.1.1 → ghostbase_cli-0.1.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ghostbase-cli"
-version = "0.1.1"
+version = "0.1.3"
 description = "A fast, keyboard-driven terminal database studio for SQLite, PostgreSQL & MySQL"
 readme = "README.md"
 authors = [{ name = "4yk5", email = "contact.4yk5@proton.me" }]

ghostbase_cli-0.1.3/tests/test_vulnerabilities_fixed.py

@@ -0,0 +1,308 @@
+import os
+import sys
+import tempfile
+import sqlite3
+import json
+import decimal
+import datetime
+import uuid
+import re
+
+# Ensure parent directory is in path
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from db_engine import DatabaseEngine
+from exporter import Exporter
+from main import detect_connection
+
+
+def test_safe_mode_and_comment_bypass():
+    print("[*] 1. Testing Safe Mode & Comment Bypass Prevention...")
+    db = DatabaseEngine(":memory:")
+
+    # Comment bypass attempt: WHERE is inside /* ... */ block comment
+    bypass_query = "UPDATE users SET role = 'Member' /* WHERE id = 1 */;"
+    is_destructive, reason = db.check_safety(bypass_query, is_safe_mode=False)
+    assert is_destructive, "VULNERABILITY: Block comment containing WHERE fooled check_safety!"
+    assert "UPDATE without WHERE" in reason
+
+    # Line comment bypass attempt: WHERE is after --
+    bypass_line = "UPDATE users SET role = 'Member' -- WHERE id = 1\n;"
+    is_destructive, reason = db.check_safety(bypass_line, is_safe_mode=False)
+    assert is_destructive, "VULNERABILITY: Line comment containing WHERE fooled check_safety!"
+
+    # String with WHERE inside must not count as WHERE clause
+    string_bypass = "UPDATE users SET role = 'WHERE -- this is a string';"
+    is_dest, r = db.check_safety(string_bypass, is_safe_mode=False)
+    assert is_dest, f"VULNERABILITY: String literal containing WHERE fooled check_safety: {r}"
+
+    # Postgres dollar-quoted string with WHERE
+    dollar_bypass = "UPDATE users SET role = $$WHERE$$;"
+    is_dest, r = db.check_safety(dollar_bypass, is_safe_mode=False)
+    assert is_dest, f"VULNERABILITY: Dollar-quoted string containing WHERE fooled check_safety: {r}"
+
+    # Postgres dollar-quoted string with tag mismatch ($tag$ vs $other$)
+    mismatched_tag_query = "UPDATE users SET role = $tag$abc $other$ WHERE marker $tag$;"
+    is_dest, r = db.check_safety(mismatched_tag_query, is_safe_mode=False)
+    assert is_dest, f"VULNERABILITY: Dollar-quoted tag mismatch bypassed WHERE check: {r}"
+
+    # Safe Mode STRICT read-only checks
+    safe_queries = [
+        "SELECT * FROM users WHERE id = 5;",
+        "SELECT 'DROP TABLE users' AS msg;",
+        "WITH cte AS (SELECT 1 AS n) SELECT * FROM cte;",
+        "EXPLAIN SELECT * FROM users;",
+    ]
+    for q in safe_queries:
+        is_dest, r = db.check_safety(q, is_safe_mode=True)
+        assert not is_dest, f"Safe query wrongly blocked in Safe Mode: {q} ({r})"
+
+    # Safe Mode must strictly block DML, DDL, DO blocks, CALL, MERGE, and SET changes
+    unsafe_queries = [
+        ("UPDATE users SET role = 'Admin' WHERE id = 5;", "UPDATE"),
+        ("INSERT INTO users (name) VALUES ('Test');", "INSERT"),
+        ("DELETE FROM users WHERE id = 5;", "DELETE"),
+        ("DROP TABLE users;", "DROP"),
+        ("TRUNCATE TABLE users;", "TRUNCATE"),
+        ("ALTER TABLE users ADD COLUMN age INT;", "ALTER"),
+        ("CREATE TABLE test (id INT);", "CREATE"),
+        ("SELECT nextval('users_id_seq');", "NEXTVAL"),
+        ("SELECT * INTO OUTFILE '/tmp/f' FROM users;", "INTO OUTFILE"),
+        ("SET default_transaction_read_only = OFF;", "SET"),
+        ("RESET ALL;", "RESET"),
+        ("DO $body$ BEGIN UPDATE users SET role = 'Admin'; END $body$;", "DO"),
+        ("CALL custom_proc();", "CALL"),
+        ("MERGE INTO target USING source ON target.id = source.id WHEN MATCHED THEN DELETE;", "MERGE"),
+    ]
+    for q, op in unsafe_queries:
+        is_dest, r = db.check_safety(q, is_safe_mode=True)
+        assert is_dest, f"VULNERABILITY: Safe Mode failed to block {op}: {q}"
+        assert "Safe Mode is active" in r
+
+    print("  [+] Safe Mode & Comment Bypass: PASS")
+
+
+def test_database_engine_level_readonly():
+    print("[*] 2. Testing Database Engine-Level Read-Only & Fail-Closed Enforcement...")
+    db = DatabaseEngine(":memory:")
+    db.conn.execute("CREATE TABLE read_only_test (id INT, val TEXT);")
+    db.conn.commit()
+
+    # Enable engine-level read-only
+    db.set_read_only(True)
+
+    # SQLite PRAGMA query_only=ON must reject writes at the C/engine level
+    failed = False
+    try:
+        db.conn.execute("INSERT INTO read_only_test VALUES (1, 'fail');")
+        db.conn.commit()
+    except sqlite3.OperationalError as e:
+        if "readonly" in str(e).lower() or "query_only" in str(e).lower():
+            failed = True
+
+    assert failed, "VULNERABILITY: SQLite allowed write with PRAGMA query_only = ON!"
+
+    # Disable read-only
+    db.set_read_only(False)
+    db.conn.execute("INSERT INTO read_only_test VALUES (1, 'success');")
+    db.conn.commit()
+
+    print("  [+] Database Engine-Level Read-Only: PASS")
+
+
+def test_sql_beautifier_preserves_literals():
+    print("[*] 3. Testing SQL Beautifier Literal Integrity (Single, Double & Dollar Quotes)...")
+    raw_query = "SELECT 'order by and where' AS message, $$order by AND where$$ AS dollar_msg FROM users WHERE id = 1;"
+    formatted = DatabaseEngine.format_sql(raw_query)
+
+    assert "'order by and where'" in formatted, f"VULNERABILITY: Beautifier corrupted single-quoted literal: {formatted}"
+    assert "$$order by AND where$$" in formatted, f"VULNERABILITY: Beautifier corrupted dollar-quoted literal: {formatted}"
+    assert "WHERE" in formatted
+    assert "SELECT" in formatted
+
+    print("  [+] SQL Beautifier Preserves Literals: PASS")
+
+
+def test_sqlite_foreign_keys_enforced():
+    print("[*] 4. Testing SQLite Foreign Keys Enforcement...")
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
+        db_path = f.name
+
+    try:
+        db = DatabaseEngine(db_path)
+        cur = db.conn.cursor()
+        cur.execute("PRAGMA foreign_keys;")
+        fk_status = cur.fetchone()[0]
+        assert fk_status == 1, f"VULNERABILITY: PRAGMA foreign_keys is OFF ({fk_status})"
+
+        db.conn.execute("CREATE TABLE dept (id INTEGER PRIMARY KEY, name TEXT);")
+        db.conn.execute("CREATE TABLE emp (id INTEGER PRIMARY KEY, dept_id INTEGER REFERENCES dept(id));")
+        db.conn.commit()
+
+        failed = False
+        try:
+            db.conn.execute("INSERT INTO emp (id, dept_id) VALUES (1, 999);")
+            db.conn.commit()
+        except sqlite3.IntegrityError:
+            failed = True
+
+        assert failed, "VULNERABILITY: SQLite allowed foreign key violation without error!"
+    finally:
+        db.close()
+        if os.path.exists(db_path):
+            os.remove(db_path)
+
+    print("  [+] SQLite Foreign Key Enforcement: PASS")
+
+
+def test_empty_sqlite_not_polluted():
+    print("[*] 5. Testing Empty SQLite Database Guard in Read-Only Mode...")
+    from app import GhostBaseApp
+    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
+        empty_db_path = f.name
+
+    try:
+        # Launch in read-only / safe mode
+        app = GhostBaseApp(db_path=empty_db_path, read_only=True, populate_sample=True)
+        tables = app.db.get_tables()
+        assert len(tables) == 0, f"VULNERABILITY: Sample tables were created during --read-only launch: {tables}"
+        app.db.close()
+    finally:
+        if os.path.exists(empty_db_path):
+            os.remove(empty_db_path)
+
+    print("  [+] Empty Database Integrity in Read-Only: PASS")
+
+
+def test_returning_dml_classification():
+    print("[*] 6. Testing DML with RETURNING Classification...")
+    db = DatabaseEngine(":memory:")
+    db.conn.execute("CREATE TABLE ret_test (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT);")
+    db.conn.commit()
+
+    res = db.execute_query("INSERT INTO ret_test (name) VALUES ('Item1') RETURNING id, name;")
+    assert res["success"]
+    assert res["is_returning"] is True, "Expected is_returning to be True"
+    assert res["is_select"] is False, "Expected is_select to be False for INSERT RETURNING"
+    assert res["affected_rows"] == 1
+    assert len(res["rows"]) == 1
+    assert res["rows"][0][1] == "Item1"
+
+    print("  [+] DML with RETURNING Classification: PASS")
+
+
+def test_exporter_json_types_and_dialect_sql():
+    print("[*] 7. Testing Exporter JSON Types & Dialect-Aware SQL Escaping (MySQL vs Postgres/SQLite)...")
+    tmp = tempfile.gettempdir()
+    json_path = os.path.join(tmp, "test_types.json")
+    sql_pg_path = os.path.join(tmp, "test_pg.sql")
+    sql_my_path = os.path.join(tmp, "test_mysql.sql")
+
+    # JSON export with Decimal, datetime, UUID, bytes, and duplicate column names
+    cols = ["id", "id", "price", "created", "uid", "raw"]
+    rows = [[
+        1, 1,
+        decimal.Decimal("49.95"),
+        datetime.datetime(2026, 11, 2, 8, 0, 0),
+        uuid.UUID("12345678-1234-5678-1234-567812345678"),
+        b"binary_ghost"
+    ]]
+    Exporter.to_json(cols, rows, json_path)
+    with open(json_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+    assert len(data) == 1
+    assert data[0]["id"] == 1
+    assert data[0]["id_2"] == 1
+    assert data[0]["price"] == 49.95
+    assert data[0]["created"] == "2026-11-02T08:00:00"
+    assert data[0]["uid"] == "12345678-1234-5678-1234-567812345678"
+    assert data[0]["raw"] == "binary_ghost"
+
+    # PostgreSQL / SQLite dialect SQL export (Double quotes)
+    table_name = "analytics.order"
+    cols_special = ["first name", "group", "price"]
+    rows_special = [["Bob", "Admin", decimal.Decimal("10.50")]]
+    Exporter.to_sql_inserts(table_name, cols_special, rows_special, sql_pg_path, dialect="postgresql")
+    with open(sql_pg_path, "r", encoding="utf-8") as f:
+        pg_content = f.read()
+
+    assert 'INSERT INTO "analytics"."order" ("first name", "group", price) VALUES (\'Bob\', \'Admin\', 10.50);' in pg_content
+
+    # MySQL dialect SQL export (Backticks)
+    Exporter.to_sql_inserts(table_name, cols_special, rows_special, sql_my_path, dialect="mysql")
+    with open(sql_my_path, "r", encoding="utf-8") as f:
+        my_content = f.read()
+
+    assert 'INSERT INTO `analytics`.`order` (`first name`, `group`, price) VALUES (\'Bob\', \'Admin\', 10.50);' in my_content
+
+    print("  [+] Exporter JSON Types & Dialect-Aware SQL Escaping: PASS")
+
+
+def test_table_quoting_and_cte_export_detection():
+    print("[*] 8. Testing Table Quoting & CTE Export Target Detection...")
+    db_sqlite = DatabaseEngine(":memory:")
+    assert db_sqlite.quote_table_name("order") == "[order]"
+
+    db_pg = DatabaseEngine.__new__(DatabaseEngine)
+    db_pg.db_type = "postgresql"
+    assert db_pg.quote_table_name("order") == '"order"'
+    assert db_pg.quote_table_name("analytics.users") == '"analytics"."users"'
+
+    db_my = DatabaseEngine.__new__(DatabaseEngine)
+    db_my.db_type = "mysql"
+    assert db_my.quote_table_name("order") == "`order`"
+    assert db_my.quote_table_name("analytics.users") == "`analytics`.`users`"
+
+    # CTE target detection test
+    cte_query = """
+    WITH recent AS (
+        SELECT * FROM archived_users WHERE active = 0
+    )
+    SELECT * FROM recent;
+    """
+    main_q = re.sub(r"^WITH\s+[\s\S]*?\)\s*", "", cte_query.strip(), flags=re.IGNORECASE)
+    m = re.search(r"\bFROM\s+([`\"\[]?[a-zA-Z0-9_\.]+(?:\"|\'|\]|`)?)", main_q, flags=re.IGNORECASE)
+    assert m is not None
+    detected = m.group(1).strip('`"[]')
+    assert detected == "recent", f"Expected target 'recent', got '{detected}'"
+
+    print("  [+] Table Quoting & CTE Export Target Detection: PASS")
+
+
+def test_result_set_memory_limit():
+    print("[*] 9. Testing Result Set Memory Limit (Capped at 1,000 rows)...")
+    db = DatabaseEngine(":memory:")
+    db.conn.execute("CREATE TABLE big (id INTEGER PRIMARY KEY);")
+    db.conn.executemany("INSERT INTO big (id) VALUES (?);", [(i,) for i in range(1500)])
+    db.conn.commit()
+
+    res = db.execute_query("SELECT * FROM big;", max_rows=1000)
+    assert res["success"]
+    assert res["row_count"] == 1000, f"Expected 1000 rows capped, got {res['row_count']}"
+    assert res["truncated"] is True, "Expected truncated flag to be True"
+
+    print("  [+] Result Set Memory Limit: PASS")
+
+
+def main():
+    print("\n" + "=" * 60)
+    print("RUNNING EXHAUSTIVE REFINED VERIFICATION TESTS (v0.1.3)")
+    print("=" * 60)
+
+    test_safe_mode_and_comment_bypass()
+    test_database_engine_level_readonly()
+    test_sql_beautifier_preserves_literals()
+    test_sqlite_foreign_keys_enforced()
+    test_empty_sqlite_not_polluted()
+    test_returning_dml_classification()
+    test_exporter_json_types_and_dialect_sql()
+    test_table_quoting_and_cte_export_detection()
+    test_result_set_memory_limit()
+
+    print("\n" + "=" * 60)
+    print("ALL ANALYST VULNERABILITIES VERIFIED AND 100% PASSING!")
+    print("=" * 60 + "\n")
+
+
+if __name__ == "__main__":
+    main()

ghostbase_cli-0.1.1/tests/test_vulnerabilities_fixed.py

@@ -1,216 +0,0 @@
-import os
-import sys
-import tempfile
-import sqlite3
-import json
-
-# Ensure parent directory is in path
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from db_engine import DatabaseEngine
-from exporter import Exporter
-from main import detect_connection
-
-
-def test_safe_mode_and_comment_bypass():
-    print("[*] 1. Testing Safe Mode & Comment Bypass Prevention...")
-    db = DatabaseEngine(":memory:")
-
-    # Comment bypass attempt: WHERE is inside /* ... */ block comment
-    bypass_query = "UPDATE users SET role = 'Member' /* WHERE id = 1 */;"
-    is_destructive, reason = db.check_safety(bypass_query, is_safe_mode=False)
-    assert is_destructive, "VULNERABILITY: Block comment containing WHERE fooled check_safety!"
-    assert "UPDATE without WHERE" in reason
-
-    # Line comment bypass attempt: WHERE is after --
-    bypass_line = "UPDATE users SET role = 'Member' -- WHERE id = 1\n;"
-    is_destructive, reason = db.check_safety(bypass_line, is_safe_mode=False)
-    assert is_destructive, "VULNERABILITY: Line comment containing WHERE fooled check_safety!"
-
-    # Safe Mode STRICT read-only checks
-    safe_queries = [
-        "SELECT * FROM users WHERE id = 5;",
-        "SELECT 'DROP TABLE users' AS msg;",
-        "WITH cte AS (SELECT 1 AS n) SELECT * FROM cte;",
-        "EXPLAIN SELECT * FROM users;",
-    ]
-    for q in safe_queries:
-        is_dest, r = db.check_safety(q, is_safe_mode=True)
-        assert not is_dest, f"Safe query wrongly blocked in Safe Mode: {q} ({r})"
-
-    unsafe_queries = [
-        ("UPDATE users SET role = 'Admin' WHERE id = 5;", "UPDATE"),
-        ("INSERT INTO users (name) VALUES ('Test');", "INSERT"),
-        ("DELETE FROM users WHERE id = 5;", "DELETE"),
-        ("DROP TABLE users;", "DROP"),
-        ("TRUNCATE TABLE users;", "TRUNCATE"),
-        ("ALTER TABLE users ADD COLUMN age INT;", "ALTER"),
-        ("CREATE TABLE test (id INT);", "CREATE"),
-    ]
-    for q, op in unsafe_queries:
-        is_dest, r = db.check_safety(q, is_safe_mode=True)
-        assert is_dest, f"VULNERABILITY: Safe Mode failed to block {op}: {q}"
-        assert "Safe Mode is active" in r
-
-    print("  [+] Safe Mode & Comment Bypass: PASS")
-
-
-def test_sql_beautifier_preserves_literals():
-    print("[*] 2. Testing SQL Beautifier Literal Integrity...")
-    raw_query = "SELECT 'order by and where' AS message, 'SELECT * FROM users' AS sample FROM users WHERE id = 1;"
-    formatted = DatabaseEngine.format_sql(raw_query)
-
-    # The string literals must not have their contents altered to uppercase or split by newlines
-    assert "'order by and where'" in formatted, f"VULNERABILITY: Beautifier corrupted string literal! Got: {formatted}"
-    assert "'SELECT * FROM users'" in formatted, "VULNERABILITY: Beautifier corrupted inner SQL string!"
-    # The actual SQL keyword outside strings should be formatted
-    assert "WHERE" in formatted
-    assert "SELECT" in formatted
-
-    print("  [+] SQL Beautifier Preserves Literals: PASS")
-
-
-def test_sqlite_foreign_keys_enforced():
-    print("[*] 3. Testing SQLite Foreign Keys Enforcement...")
-    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
-        db_path = f.name
-
-    try:
-        db = DatabaseEngine(db_path)
-        # Verify PRAGMA foreign_keys is ON
-        cur = db.conn.cursor()
-        cur.execute("PRAGMA foreign_keys;")
-        fk_status = cur.fetchone()[0]
-        assert fk_status == 1, f"VULNERABILITY: PRAGMA foreign_keys is OFF ({fk_status})"
-
-        # Verify real constraint enforcement
-        db.conn.execute("CREATE TABLE dept (id INTEGER PRIMARY KEY, name TEXT);")
-        db.conn.execute("CREATE TABLE emp (id INTEGER PRIMARY KEY, dept_id INTEGER REFERENCES dept(id));")
-        db.conn.commit()
-
-        # Inserting child with invalid foreign key must raise IntegrityError
-        failed = False
-        try:
-            db.conn.execute("INSERT INTO emp (id, dept_id) VALUES (1, 999);")
-            db.conn.commit()
-        except sqlite3.IntegrityError:
-            failed = True
-
-        assert failed, "VULNERABILITY: SQLite allowed foreign key violation without error!"
-    finally:
-        db.close()
-        if os.path.exists(db_path):
-            os.remove(db_path)
-
-    print("  [+] SQLite Foreign Key Enforcement: PASS")
-
-
-def test_empty_sqlite_not_polluted():
-    print("[*] 4. Testing Empty SQLite Database Guard...")
-    from app import GhostBaseApp
-    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
-        empty_db_path = f.name
-
-    try:
-        # User opens an empty database without populate_sample flag
-        app = GhostBaseApp(db_path=empty_db_path, populate_sample=False)
-        tables = app.db.get_tables()
-        assert len(tables) == 0, f"VULNERABILITY: Sample tables were automatically added to user's database: {tables}"
-        app.db.close()
-    finally:
-        if os.path.exists(empty_db_path):
-            os.remove(empty_db_path)
-
-    print("  [+] Empty Database Integrity: PASS")
-
-
-def test_exporter_dedup_and_ident_escaping():
-    print("[*] 5. Testing Exporter Duplicate Column Dedup & Quoting...")
-    tmp = tempfile.gettempdir()
-    json_path = os.path.join(tmp, "test_dedup.json")
-    sql_path = os.path.join(tmp, "test_escape.sql")
-
-    # JSON duplicate columns
-    cols = ["id", "id", "name"]
-    rows = [[1, 2, "Alice"]]
-    Exporter.to_json(cols, rows, json_path)
-    with open(json_path, "r", encoding="utf-8") as f:
-        data = json.load(f)
-    assert len(data) == 1
-    assert data[0]["id"] == 1
-    assert data[0]["id_2"] == 2
-    assert data[0]["name"] == "Alice"
-
-    # SQL export reserved keywords & special characters
-    table_name = "order"  # Reserved word
-    cols_with_special = ["first name", "group", "price"]  # "first name" has space, "group" is reserved
-    rows_special = [[ "Bob", "Admin", 99.5 ]]
-    Exporter.to_sql_inserts(table_name, cols_with_special, rows_special, sql_path)
-    with open(sql_path, "r", encoding="utf-8") as f:
-        sql_content = f.read()
-
-    assert 'INSERT INTO "order" ("first name", "group", price)' in sql_content
-
-    print("  [+] Exporter Dedup & Escaping: PASS")
-
-
-def test_result_set_memory_limit():
-    print("[*] 6. Testing Result Set Memory Limit (Capped at 1,000 rows)...")
-    db = DatabaseEngine(":memory:")
-    db.conn.execute("CREATE TABLE big (id INTEGER PRIMARY KEY);")
-    db.conn.executemany("INSERT INTO big (id) VALUES (?);", [(i,) for i in range(1500)])
-    db.conn.commit()
-
-    res = db.execute_query("SELECT * FROM big;", max_rows=1000)
-    assert res["success"]
-    assert res["row_count"] == 1000, f"Expected 1000 rows capped, got {res['row_count']}"
-    assert res["truncated"] is True, "Expected truncated flag to be True"
-
-    print("  [+] Result Set Memory Limit: PASS")
-
-
-def test_env_credentials_url_encoding():
-    print("[*] 7. Testing .env Credentials URL Encoding...")
-    tmp_env = os.path.join(tempfile.gettempdir(), ".env")
-    with open(tmp_env, "w", encoding="utf-8") as f:
-        f.write("DB_HOST=127.0.0.1\n")
-        f.write("DB_PORT=5432\n")
-        f.write("DB_NAME=my_db\n")
-        f.write("DB_USER=user@work#1\n")
-        f.write("DB_PASSWORD=p@ss:w/ord%!\n")
-
-    # Backup current dir and test detect_connection
-    orig_dir = os.getcwd()
-    os.chdir(tempfile.gettempdir())
-    try:
-        url = detect_connection()
-        assert "user%40work%231" in url
-        assert "p%40ss%3Aw%2Ford%25%21" in url
-    finally:
-        os.chdir(orig_dir)
-        if os.path.exists(tmp_env):
-            os.remove(tmp_env)
-
-    print("  [+] .env Credentials URL Encoding: PASS")
-
-
-def main():
-    print("\n" + "=" * 60)
-    print("RUNNING SECURITY & BUG FIX VERIFICATION TESTS")
-    print("=" * 60)
-
-    test_safe_mode_and_comment_bypass()
-    test_sql_beautifier_preserves_literals()
-    test_sqlite_foreign_keys_enforced()
-    test_empty_sqlite_not_polluted()
-    test_exporter_dedup_and_ident_escaping()
-    test_result_set_memory_limit()
-    test_env_credentials_url_encoding()
-
-    print("\n" + "=" * 60)
-    print("ALL 10 ANALYST FINDINGS & VULNERABILITIES ARE 100% FIXED!")
-    print("=" * 60 + "\n")
-
-
-if __name__ == "__main__":
-    main()