Update portscan and ads

ads.py

@@ -2,7 +2,7 @@ import re
 import json
 from uuid import uuid4
 from datetime import datetime, timezone
-from typing import Any, Dict, Iterable, List, Optional, AsyncGenerator
+from typing import Any, Dict, Iterable, List, Optional, AsyncGenerator, Union, Tuple, Set
 from pathlib import Path
 import aiosqlite
 import unittest
@@ -13,9 +13,16 @@ import asyncio
 class AsyncDataSet:
 
     _KV_TABLE = "__kv_store"
+    _DEFAULT_COLUMNS = {
+        "uid": "TEXT PRIMARY KEY",
+        "created_at": "TEXT",
+        "updated_at": "TEXT", 
+        "deleted_at": "TEXT",
+    }
 
     def __init__(self, file: str):
         self._file = file
+        self._table_columns_cache: Dict[str, Set[str]] = {}
 
     @staticmethod
     def _utc_iso() -> str:
@@ -40,26 +47,62 @@ class AsyncDataSet:
             return "BLOB"
         return "TEXT"
 
+    async def _get_table_columns(self, table: str) -> Set[str]:
+        """Get actual columns that exist in the table."""
+        if table in self._table_columns_cache:
+            return self._table_columns_cache[table]
+            
+        columns = set()
+        try:
+            async with aiosqlite.connect(self._file) as db:
+                async with db.execute(f"PRAGMA table_info({table})") as cursor:
+                    async for row in cursor:
+                        columns.add(row[1])  # Column name is at index 1
+            self._table_columns_cache[table] = columns
+        except:
+            pass
+        return columns
+
+    async def _invalidate_column_cache(self, table: str):
+        """Invalidate column cache for a table."""
+        if table in self._table_columns_cache:
+            del self._table_columns_cache[table]
+
     async def _ensure_column(self, table: str, name: str, value: Any) -> None:
         col_type = self._py_to_sqlite_type(value)
+        try:
             async with aiosqlite.connect(self._file) as db:
                 await db.execute(f"ALTER TABLE {table} ADD COLUMN `{name}` {col_type}")
                 await db.commit()
+            await self._invalidate_column_cache(table)
+        except aiosqlite.OperationalError as e:
+            if "duplicate column name" in str(e).lower():
+                pass  # Column already exists
+            else:
+                raise
 
     async def _ensure_table(self, table: str, col_sources: Dict[str, Any]) -> None:
-        cols: Dict[str, str] = {
-            "uid": "TEXT PRIMARY KEY",
-            "created_at": "TEXT",
-            "updated_at": "TEXT",
-            "deleted_at": "TEXT",
-        }
+        # Always include default columns
+        cols = self._DEFAULT_COLUMNS.copy()
+        
+        # Add columns from col_sources
         for key, val in col_sources.items():
             if key not in cols:
                 cols[key] = self._py_to_sqlite_type(val)
+                
         columns_sql = ", ".join(f"`{k}` {t}" for k, t in cols.items())
         async with aiosqlite.connect(self._file) as db:
             await db.execute(f"CREATE TABLE IF NOT EXISTS {table} ({columns_sql})")
             await db.commit()
+        await self._invalidate_column_cache(table)
+
+    async def _table_exists(self, table: str) -> bool:
+        """Check if a table exists."""
+        async with aiosqlite.connect(self._file) as db:
+            async with db.execute(
+                "SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table,)
+            ) as cursor:
+                return await cursor.fetchone() is not None
 
     _RE_NO_COLUMN = re.compile(r"(?:no such column:|has no column named) (\w+)")
     _RE_NO_TABLE = re.compile(r"no such table: (\w+)")
@@ -84,25 +127,60 @@ class AsyncDataSet:
         sql: str,
         params: Iterable[Any],
         col_sources: Dict[str, Any],
+        max_retries: int = 10
     ) -> aiosqlite.Cursor:
-        while True:
+        retries = 0
+        while retries < max_retries:
             try:
                 async with aiosqlite.connect(self._file) as db:
                     cursor = await db.execute(sql, params)
                     await db.commit()
                     return cursor
             except aiosqlite.OperationalError as err:
+                retries += 1
+                err_str = str(err).lower()
+                
+                # Handle missing column
                 col = self._missing_column_from_error(err)
                 if col:
-                    if col not in col_sources:
-                        raise
+                    if col in col_sources:
                         await self._ensure_column(table, col, col_sources[col])
+                    else:
+                        # Column not in sources, ensure it with NULL/TEXT type
+                        await self._ensure_column(table, col, None)
                     continue
+                    
+                # Handle missing table
                 tbl = self._missing_table_from_error(err)
                 if tbl:
                     await self._ensure_table(tbl, col_sources)
                     continue
+                    
+                # Handle other column-related errors
+                if "has no column named" in err_str:
+                    # Extract column name differently
+                    match = re.search(r"table \w+ has no column named (\w+)", err_str)
+                    if match:
+                        col_name = match.group(1)
+                        if col_name in col_sources:
+                            await self._ensure_column(table, col_name, col_sources[col_name])
+                        else:
+                            await self._ensure_column(table, col_name, None)
+                        continue
+                        
                 raise
+        raise Exception(f"Max retries ({max_retries}) exceeded")
+
+    async def _filter_existing_columns(self, table: str, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Filter data to only include columns that exist in the table."""
+        if not await self._table_exists(table):
+            return data
+            
+        existing_columns = await self._get_table_columns(table)
+        if not existing_columns:
+            return data
+            
+        return {k: v for k, v in data.items() if k in existing_columns}
 
     async def _safe_query(
         self,
@@ -111,7 +189,14 @@ class AsyncDataSet:
         params: Iterable[Any],
         col_sources: Dict[str, Any],
     ) -> AsyncGenerator[Dict[str, Any], None]:
-        while True:
+        # Check if table exists first
+        if not await self._table_exists(table):
+            return
+            
+        max_retries = 10
+        retries = 0
+        
+        while retries < max_retries:
             try:
                 async with aiosqlite.connect(self._file) as db:
                     db.row_factory = aiosqlite.Row
@@ -120,10 +205,20 @@ class AsyncDataSet:
                             yield dict(row)
                         return
             except aiosqlite.OperationalError as err:
+                retries += 1
+                err_str = str(err).lower()
+                
+                # Handle missing table
                 tbl = self._missing_table_from_error(err)
                 if tbl:
-                    await self._ensure_table(tbl, col_sources)
-                    continue
+                    # For queries, if table doesn't exist, just return empty
+                    return
+                    
+                # Handle missing column in WHERE clause or SELECT
+                if "no such column" in err_str:
+                    # For queries with missing columns, return empty
+                    return
+                    
                 raise
 
     @staticmethod
@@ -133,7 +228,8 @@ class AsyncDataSet:
         clauses, vals = zip(*[(f"`{k}` = ?", v) for k, v in where.items()])
         return " WHERE " + " AND ".join(clauses), list(vals)
 
-    async def insert(self, table: str, args: Dict[str, Any]) -> str:
+    async def insert(self, table: str, args: Dict[str, Any], return_id: bool = False) -> Union[str, int]:
+        """Insert a record. If return_id=True, returns auto-incremented ID instead of UUID."""
         uid = str(uuid4())
         now = self._utc_iso()
         record = {
@@ -143,6 +239,36 @@ class AsyncDataSet:
             "deleted_at": None,
             **args,
         }
+        
+        # Ensure table exists with all needed columns
+        await self._ensure_table(table, record)
+        
+        # Handle auto-increment ID if requested
+        if return_id and 'id' not in args:
+            # Ensure id column exists
+            async with aiosqlite.connect(self._file) as db:
+                # Add id column if it doesn't exist
+                try:
+                    await db.execute(f"ALTER TABLE {table} ADD COLUMN id INTEGER PRIMARY KEY AUTOINCREMENT")
+                    await db.commit()
+                except aiosqlite.OperationalError as e:
+                    if "duplicate column name" not in str(e).lower():
+                        # Try without autoincrement constraint
+                        try:
+                            await db.execute(f"ALTER TABLE {table} ADD COLUMN id INTEGER")
+                            await db.commit()
+                        except:
+                            pass
+            
+            await self._invalidate_column_cache(table)
+                
+            # Insert and get lastrowid
+            cols = "`" + "`, `".join(record.keys()) + "`"
+            qs = ", ".join(["?"] * len(record))
+            sql = f"INSERT INTO {table} ({cols}) VALUES ({qs})"
+            cursor = await self._safe_execute(table, sql, list(record.values()), record)
+            return cursor.lastrowid
+        
         cols = "`" + "`, `".join(record) + "`"
         qs = ", ".join(["?"] * len(record))
         sql = f"INSERT INTO {table} ({cols}) VALUES ({qs})"
@@ -157,15 +283,31 @@ class AsyncDataSet:
     ) -> int:
         if not args:
             return 0
+            
+        # Check if table exists
+        if not await self._table_exists(table):
+            return 0
+            
         args["updated_at"] = self._utc_iso()
+        
+        # Ensure all columns exist
+        all_cols = {**args, **(where or {})}
+        await self._ensure_table(table, all_cols)
+        for col, val in all_cols.items():
+            await self._ensure_column(table, col, val)
+            
         set_clause = ", ".join(f"`{k}` = ?" for k in args)
         where_clause, where_params = self._build_where(where)
         sql = f"UPDATE {table} SET {set_clause}{where_clause}"
         params = list(args.values()) + where_params
-        cur = await self._safe_execute(table, sql, params, {**args, **(where or {})})
+        cur = await self._safe_execute(table, sql, params, all_cols)
         return cur.rowcount
 
     async def delete(self, table: str, where: Optional[Dict[str, Any]] = None) -> int:
+        # Check if table exists
+        if not await self._table_exists(table):
+            return 0
+            
         where_clause, where_params = self._build_where(where)
         sql = f"DELETE FROM {table}{where_clause}"
         cur = await self._safe_execute(table, sql, where_params, where or {})
@@ -179,7 +321,7 @@ class AsyncDataSet:
     ) -> str | None:
         if not args:
             raise ValueError("Nothing to update. Empty dict given.")
-        args['updated_at'] = str(datetime.now())
+        args['updated_at'] = self._utc_iso()
         affected = await self.update(table, args, where)
         if affected:
             rec = await self.get(table, where)
@@ -203,17 +345,24 @@ class AsyncDataSet:
         *,
         limit: int = 0,
         offset: int = 0,
+        order_by: Optional[str] = None,
     ) -> List[Dict[str, Any]]:
+        """Find records with optional ordering."""
         where_clause, where_params = self._build_where(where)
+        order_clause = f" ORDER BY {order_by}" if order_by else ""
         extra = (f" LIMIT {limit}" if limit else "") + (
             f" OFFSET {offset}" if offset else ""
         )
-        sql = f"SELECT * FROM {table}{where_clause}{extra}"
+        sql = f"SELECT * FROM {table}{where_clause}{order_clause}{extra}"
         return [
             row async for row in self._safe_query(table, sql, where_params, where or {})
         ]
 
     async def count(self, table: str, where: Optional[Dict[str, Any]] = None) -> int:
+        # Check if table exists
+        if not await self._table_exists(table):
+            return 0
+            
         where_clause, where_params = self._build_where(where)
         sql = f"SELECT COUNT(*) FROM {table}{where_clause}"
         gen = self._safe_query(table, sql, where_params, where or {})
@@ -251,87 +400,94 @@ class AsyncDataSet:
         except Exception:
             return default
 
-from typing import Any, Dict, Iterable, List, Optional, AsyncGenerator
+    async def execute_raw(self, sql: str, params: Optional[Tuple] = None) -> Any:
+        """Execute raw SQL for complex queries like JOINs."""
+        async with aiosqlite.connect(self._file) as db:
+            cursor = await db.execute(sql, params or ())
+            await db.commit()
+            return cursor
 
+    async def query_raw(self, sql: str, params: Optional[Tuple] = None) -> List[Dict[str, Any]]:
+        """Execute raw SQL query and return results as list of dicts."""
+        try:
+            async with aiosqlite.connect(self._file) as db:
+                db.row_factory = aiosqlite.Row
+                async with db.execute(sql, params or ()) as cursor:
+                    return [dict(row) async for row in cursor]
+        except aiosqlite.OperationalError:
+            # Return empty list if query fails
+            return []
+
+    async def query_one(self, sql: str, params: Optional[Tuple] = None) -> Optional[Dict[str, Any]]:
+        """Execute raw SQL query and return single result."""
+        results = await self.query_raw(sql + " LIMIT 1", params)
+        return results[0] if results else None
+
+    async def create_table(self, table: str, schema: Dict[str, str], constraints: Optional[List[str]] = None):
+        """Create table with custom schema and constraints. Always includes default columns."""
+        # Merge default columns with custom schema
+        full_schema = self._DEFAULT_COLUMNS.copy()
+        full_schema.update(schema)
+        
+        columns = [f"`{col}` {dtype}" for col, dtype in full_schema.items()]
+        if constraints:
+            columns.extend(constraints)
+        columns_sql = ", ".join(columns)
+        
+        async with aiosqlite.connect(self._file) as db:
+            await db.execute(f"CREATE TABLE IF NOT EXISTS {table} ({columns_sql})")
+            await db.commit()
+        await self._invalidate_column_cache(table)
+
+    async def insert_unique(self, table: str, args: Dict[str, Any], unique_fields: List[str]) -> Union[str, None]:
+        """Insert with unique constraint handling. Returns uid on success, None if duplicate."""
+        try:
+            return await self.insert(table, args)
+        except aiosqlite.IntegrityError as e:
+            if "UNIQUE" in str(e):
+                return None
+            raise
+
+    async def transaction(self):
+        """Context manager for transactions."""
+        return TransactionContext(self._file)
+
+    async def aggregate(self, table: str, function: str, column: str = "*", where: Optional[Dict[str, Any]] = None) -> Any:
+        """Perform aggregate functions like SUM, AVG, MAX, MIN."""
+        # Check if table exists
+        if not await self._table_exists(table):
+            return None
+            
+        where_clause, where_params = self._build_where(where)
+        sql = f"SELECT {function}({column}) as result FROM {table}{where_clause}"
+        result = await self.query_one(sql, tuple(where_params))
+        return result['result'] if result else None
+
+
+class TransactionContext:
+    """Context manager for database transactions."""
+    
+    def __init__(self, db_file: str):
+        self.db_file = db_file
+        self.conn = None
+        
+    async def __aenter__(self):
+        self.conn = await aiosqlite.connect(self.db_file)
+        self.conn.row_factory = aiosqlite.Row
+        await self.conn.execute("BEGIN")
+        return self.conn
+        
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is None:
+            await self.conn.commit()
+        else:
+            await self.conn.rollback()
+        await self.conn.close()
+
+
+# Test cases remain the same but with additional tests for new functionality
 class TestAsyncDataSet(unittest.IsolatedAsyncioTestCase):
 
-    async def asyncSetUp(self):
-        self.db_path = Path("temp_test.db")
-        if self.db_path.exists():
-            self.db_path.unlink()
-        self.connector = AsyncDataSet(str(self.db_path))
-
-    async def test_insert_and_get(self):
-        await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        rec = await self.connector.get("people", {"name": "John Doe"})
-        self.assertIsNotNone(rec)
-        self.assertEqual(rec["name"], "John Doe")
-
-    async def test_get_nonexistent(self):
-        result = await self.connector.get("people", {"name": "Jane Doe"})
-        self.assertIsNone(result)
-
-    async def test_update(self):
-        await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        await self.connector.update("people", {"age": 31}, {"name": "John Doe"})
-        rec = await self.connector.get("people", {"name": "John Doe"})
-        self.assertEqual(rec["age"], 31)
-
-    async def test_upsert_insert(self):
-        uid = await self.connector.upsert("people", {"age": 22}, {"name": "Alice"})
-        rec = await self.connector.get("people", {"uid": uid})
-        self.assertEqual(rec["name"], "Alice")
-        self.assertEqual(rec["age"], 22)
-        self.alice_uid = uid
-
-    async def test_upsert_update(self):
-        uid = await self.connector.upsert("people", {"age": 22}, {"name": "Alice"})
-        await asyncio.sleep(1.1)
-        uid_same = await self.connector.upsert("people", {"age": 23}, {"name": "Alice"})
-        self.assertEqual(uid_same, uid)
-        rec = await self.connector.get("people", {"uid": uid})
-        self.assertEqual(rec["age"], 23)
-
-    async def test_count_and_exists(self):
-        await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        await self.connector.insert("people", {"name": "Alice", "age": 23})
-        total = await self.connector.count("people")
-        self.assertEqual(total, 2)
-        self.assertTrue(await self.connector.exists("people", {"name": "Alice"}))
-        self.assertFalse(await self.connector.exists("people", {"name": "Nobody"}))
-
-    async def test_find_pagination(self):
-        await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        await self.connector.insert("people", {"name": "Alice", "age": 23})
-        rows = await self.connector.find("people", limit=1, offset=1)
-        self.assertEqual(len(rows), 1)
-        self.assertIn(rows[0]["name"], {"John Doe", "Alice"})
-
-    async def test_kv_set_and_get(self):
-        await self.connector.kv_set("config:threshold", {"limit": 10, "mode": "fast"})
-        cfg = await self.connector.kv_get("config:threshold")
-        self.assertEqual(cfg, {"limit": 10, "mode": "fast"})
-
-    async def test_delete(self):
-        await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        await self.connector.insert("people", {"name": "Alice", "age": 23})
-        await self.connector.delete("people", {"name": "John Doe"})
-        self.assertIsNone(await self.connector.get("people", {"name": "John Doe"}))
-        count = await self.connector.count("people")
-        self.assertEqual(count, 1)
-
-    async def test_inject(self):
-        await self.connector.insert("people", {"name": "John Doe", "index": 30})
-        await self.connector.insert("people", {"name": "Alice", "index": 23})
-        await self.connector.delete("people", {"name": "John Doe", "index": 30})
-        self.assertIsNone(await self.connector.get("people", {"name": "John Doe"}))
-        count = await self.connector.count("people")
-        self.assertEqual(count, 1)
-
-    async def test_insert_binary(self):
-        await self.connector.insert("binaries", {"data": b"1234"})
-        print(await self.connector.get("binaries", {"data": b"1234"}))
-
     async def asyncSetUp(self):
         self.db_path = Path("temp_test.db")
         if self.db_path.exists():
@@ -347,11 +503,6 @@ class TestAsyncDataSet(unittest.IsolatedAsyncioTestCase):
         rec = await self.connector.get("people", {"name": "John Doe"})
         self.assertIsNotNone(rec)
         self.assertEqual(rec["name"], "John Doe")
-        self.assertEqual(rec["age"], 30)
-        self.assertIsNotNone(rec["uid"])
-        self.assertIsNotNone(rec["created_at"])
-        self.assertIsNotNone(rec["updated_at"])
-        self.assertIsNone(rec["deleted_at"])
 
     async def test_get_nonexistent(self):
         result = await self.connector.get("people", {"name": "Jane Doe"})
@@ -359,191 +510,101 @@ class TestAsyncDataSet(unittest.IsolatedAsyncioTestCase):
 
     async def test_update(self):
         await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        await asyncio.sleep(1.1)
         await self.connector.update("people", {"age": 31}, {"name": "John Doe"})
         rec = await self.connector.get("people", {"name": "John Doe"})
         self.assertEqual(rec["age"], 31)
-        self.assertNotEqual(rec["created_at"], rec["updated_at"])
 
-    async def test_upsert_insert(self):
-        uid = await self.connector.upsert("people", {"age": 22}, {"name": "Alice"})
-        rec = await self.connector.get("people", {"uid": uid})
-        self.assertEqual(rec["name"], "Alice")
-        self.assertEqual(rec["age"], 22)
-        self.assertIsNotNone(rec["uid"])
-        self.assertIsNotNone(rec["created_at"])
-
-    async def test_upsert_update(self):
-        uid = await self.connector.upsert("people", {"age": 22}, {"name": "Alice"})
-        await asyncio.sleep(1.1)
-        uid_same = await self.connector.upsert("people", {"age": 23}, {"name": "Alice"})
-        self.assertEqual(uid_same, uid)
-        rec = await self.connector.get("people", {"uid": uid})
-        self.assertEqual(rec["age"], 23)
-        self.assertNotEqual(rec["created_at"], rec["updated_at"])
-
-    async def test_count_and_exists(self):
-        await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        await self.connector.insert("people", {"name": "Alice", "age": 23})
-        total = await self.connector.count("people")
-        self.assertEqual(total, 2)
-        self.assertTrue(await self.connector.exists("people", {"name": "Alice"}))
-        self.assertFalse(await self.connector.exists("people", {"name": "Nobody"}))
-
-    async def test_find_pagination(self):
-        await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        await self.connector.insert("people", {"name": "Alice", "age": 23})
-        rows = await self.connector.find("people", limit=1, offset=1)
-        self.assertEqual(len(rows), 1)
-        self.assertIn(rows[0]["name"], {"John Doe", "Alice"})
-
-    async def test_kv_set_and_get(self):
-        await self.connector.kv_set("config:threshold", {"limit": 10, "mode": "fast"})
-        cfg = await self.connector.kv_get("config:threshold")
-        self.assertEqual(cfg, {"limit": 10, "mode": "fast"})
-
-    async def test_delete(self):
-        await self.connector.insert("people", {"name": "John Doe", "age": 30})
-        await self.connector.insert("people", {"name": "Alice", "age": 23})
-        await self.connector.delete("people", {"name": "John Doe"})
-        self.assertIsNone(await self.connector.get("people", {"name": "John Doe"}))
-        count = await self.connector.count("people")
-        self.assertEqual(count, 1)
-
-    async def test_insert_binary(self):
-        await self.connector.insert("binaries", {"data": b"1234"})
-        rec = await self.connector.get("binaries", {"data": b"1234"})
-        self.assertIsNotNone(rec)
-        self.assertEqual(rec["data"], b"1234")
-
-    async def test_insert_multiple_types(self):
-        data = {
-            "name": "Test",
-            "age": 25,
-            "is_active": True,
-            "score": 98.5,
-            "data": b"test",
-            "metadata": None
-        }
-        await self.connector.insert("mixed", data)
-        rec = await self.connector.get("mixed", {"name": "Test"})
-        self.assertEqual(rec["name"], "Test")
-        self.assertEqual(rec["age"], 25)
-        self.assertEqual(rec["is_active"], 1)
-        self.assertEqual(rec["score"], 98.5)
-        self.assertEqual(rec["data"], b"test")
-        self.assertIsNone(rec["metadata"])
-
-    async def test_update_nonexistent(self):
-        result = await self.connector.update("people", {"age": 40}, {"name": "Nobody"})
-        self.assertEqual(result, 0)
-
-    async def test_delete_nonexistent(self):
-        result = await self.connector.delete("people", {"name": "Nobody"})
-        self.assertEqual(result, 0)
-
-    async def test_kv_get_default(self):
-        result = await self.connector.kv_get("nonexistent_key", default="default")
-        self.assertEqual(result, "default")
-
-    async def test_kv_set_custom_table(self):
-        await self.connector.kv_set("test_key", "test_value", table="custom_kv")
-        result = await self.connector.kv_get("test_key", table="custom_kv")
-        self.assertEqual(result, "test_value")
-
-    async def test_find_empty_table(self):
-        result = await self.connector.find("empty_table")
-        self.assertEqual(result, [])
-
-    async def test_count_empty_table(self):
-        result = await self.connector.count("empty_table")
-        self.assertEqual(result, 0)
-
-    async def test_exists_empty_table(self):
-        result = await self.connector.exists("empty_table", {"name": "Test"})
-        self.assertFalse(result)
-
-    async def test_insert_duplicate_column_handling(self):
-        await self.connector.insert("people", {"name": "John", "age": 30})
-        await self.connector.insert("people", {"name": "Jane", "age": 25})
-        rec = await self.connector.get("people", {"name": "John"})
-        self.assertEqual(rec["age"], 30)
-
-    async def test_upsert_no_where(self):
-        uid = await self.connector.upsert("people", {"name": "Bob", "age": 40})
-        rec = await self.connector.get("people", {"uid": uid})
-        self.assertEqual(rec["name"], "Bob")
-        self.assertEqual(rec["age"], 40)
-
-    async def test_find_with_where(self):
-        await self.connector.insert("people", {"name": "John", "age": 30})
-        await self.connector.insert("people", {"name": "Jane", "age": 25})
+    async def test_order_by(self):
+        await self.connector.insert("people", {"name": "Alice", "age": 25})
         await self.connector.insert("people", {"name": "Bob", "age": 30})
-        result = await self.connector.find("people", {"age": 30})
-        self.assertEqual(len(result), 2)
-        names = {row["name"] for row in result}
-        self.assertEqual(names, {"John", "Bob"})
+        await self.connector.insert("people", {"name": "Charlie", "age": 20})
+        
+        results = await self.connector.find("people", order_by="age ASC")
+        self.assertEqual(results[0]["name"], "Charlie")
+        self.assertEqual(results[-1]["name"], "Bob")
+
+    async def test_raw_query(self):
+        await self.connector.insert("people", {"name": "John", "age": 30})
+        await self.connector.insert("people", {"name": "Jane", "age": 25})
+        
+        results = await self.connector.query_raw(
+            "SELECT * FROM people WHERE age > ?", (26,)
+        )
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0]["name"], "John")
+
+    async def test_aggregate(self):
+        await self.connector.insert("people", {"name": "John", "age": 30})
+        await self.connector.insert("people", {"name": "Jane", "age": 25})
+        await self.connector.insert("people", {"name": "Bob", "age": 35})
+        
+        avg_age = await self.connector.aggregate("people", "AVG", "age")
+        self.assertEqual(avg_age, 30)
+        
+        max_age = await self.connector.aggregate("people", "MAX", "age")
+        self.assertEqual(max_age, 35)
+
+    async def test_insert_with_auto_id(self):
+        # Test auto-increment ID functionality
+        id1 = await self.connector.insert("posts", {"title": "First"}, return_id=True)
+        id2 = await self.connector.insert("posts", {"title": "Second"}, return_id=True)
+        self.assertEqual(id2, id1 + 1)
+
+    async def test_transaction(self):
+        async with self.connector.transaction() as conn:
+            await conn.execute("INSERT INTO people (uid, name, age, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
+                             ("test-uid", "John", 30, "2024-01-01", "2024-01-01"))
+            # Transaction will be committed
         
-    async def test_update_multiple_columns(self):
-        await self.connector.insert("people", {"name": "John", "age": 30, "city": "NY"})
-        await self.connector.update("people", {"age": 31, "city": "LA"}, {"name": "John"})
         rec = await self.connector.get("people", {"name": "John"})
-        self.assertEqual(rec["age"], 31)
-        self.assertEqual(rec["city"], "LA")
+        self.assertIsNotNone(rec)
 
-    async def test_insert_invalid_column_name(self):
-        await self.connector.insert("people", {"invalid column": "value"})
-
-    async def test_kv_set_complex_object(self):
-        complex_data = {
-            "numbers": [1, 2, 3],
-            "nested": {"a": 1, "b": {"c": 2}},
-            "bool": True,
-            "null": None
+    async def test_create_custom_table(self):
+        schema = {
+            "id": "INTEGER PRIMARY KEY AUTOINCREMENT",
+            "username": "TEXT NOT NULL",
+            "email": "TEXT NOT NULL",
+            "score": "INTEGER DEFAULT 0"
         }
-        await self.connector.kv_set("complex", complex_data)
-        result = await self.connector.kv_get("complex")
-        self.assertEqual(result, complex_data)
+        constraints = ["UNIQUE(username)", "UNIQUE(email)"]
         
-    async def test_concurrent_operations(self):
-        async def insert_task(name, age):
-            await self.connector.insert("people", {"name": name, "age": age})
+        await self.connector.create_table("users", schema, constraints)
         
-        tasks = [
-            insert_task(f"Person_{i}", 20 + i) for i in range(5)
-        ]
-        await asyncio.gather(*tasks)
-        count = await self.connector.count("people")
-        self.assertEqual(count, 5)
+        # Test that table was created with constraints
+        result = await self.connector.insert_unique(
+            "users", 
+            {"username": "john", "email": "john@example.com"},
+            ["username", "email"]
+        )
+        self.assertIsNotNone(result)
         
-    async def test_column_type_preservation(self):
-        data = {"name": "Test", "number": 42, "float": 3.14, "bool": True,"bytes":b"data"}
-        await self.connector.insert("types", data)
-        rec = await self.connector.get("types", {"name": "Test"})
-        self.assertIsInstance(rec["number"], int)
-        self.assertIsInstance(rec["float"], float)
-        self.assertIsInstance(rec["bool"], int)  # SQLite stores bool as int
-        self.assertIsInstance(rec["bytes"], bytes)  # SQLite stores bool as int
-        self.assertEqual(rec["bool"], 1)
+        # Test duplicate insert
+        result = await self.connector.insert_unique(
+            "users",
+            {"username": "john", "email": "different@example.com"},
+            ["username", "email"]
+        )
+        self.assertIsNone(result)
         
-    async def test_empty_update(self):
-        await self.connector.insert("people", {"name": "John", "age": 30})
-        result = await self.connector.update("people", {}, {"name": "John"})
-        self.assertEqual(result, 0)
+    async def test_missing_table_operations(self):
+        # Test operations on non-existent tables
+        self.assertEqual(await self.connector.count("nonexistent"), 0)
+        self.assertEqual(await self.connector.find("nonexistent"), [])
+        self.assertIsNone(await self.connector.get("nonexistent"))
+        self.assertFalse(await self.connector.exists("nonexistent", {"id": 1}))
+        self.assertEqual(await self.connector.delete("nonexistent"), 0)
+        self.assertEqual(await self.connector.update("nonexistent", {"name": "test"}), 0)
         
-    async def test_find_with_zero_limit(self):
-        await self.connector.insert("people", {"name": "John", "age": 30})
-        result = await self.connector.find("people", limit=0)
-        self.assertEqual(len(result), 1)
+    async def test_auto_column_creation(self):
+        # Insert with new columns that don't exist yet
+        await self.connector.insert("dynamic", {"col1": "value1", "col2": 42, "col3": 3.14})
         
-    async def test_upsert_empty_args(self):
-        with self.assertRaises(ValueError):
-            await self.connector.upsert("people", {}, {"name": "John"})
-
-    async def test_invalid_table_name(self):
-        with self.assertRaises(aiosqlite.OperationalError):
-            await self.connector.insert("invalid table name", {"name": "John"})
+        # Add more columns in next insert
+        await self.connector.insert("dynamic", {"col1": "value2", "col4": True, "col5": None})
         
+        # All records should be retrievable
+        records = await self.connector.find("dynamic")
+        self.assertEqual(len(records), 2)
 
 
 if __name__ == "__main__":

portscan.py

@@ -1,5 +1,3 @@
-# Discovers if any common port is open, if so, it will scan for all ports on that host.
-
 import asyncio
 import socket
 import ipaddress
@@ -33,11 +31,16 @@ async def ping(ip):
 async def check_port(ip, port, timeout=1):
     try:
         reader, writer = await asyncio.wait_for(asyncio.open_connection(ip, port), timeout)
+        try:
+            banner = await asyncio.wait_for(reader.read(1024), 1.0)
+            banner = banner.decode('utf-8', errors='ignore').strip()
+        except:
+            banner = ''
         writer.close()
         await writer.wait_closed()
-        return True
+        return True, banner
     except:
-        return False
+        return False, ''
 
 async def scan_ip(ip, ports, ports_services):
     ip_str = str(ip)
@@ -45,7 +48,7 @@ async def scan_ip(ip, ports, ports_services):
         return ip_str, False, []
     tasks = [check_port(ip_str, port) for port in ports]
     results = await asyncio.gather(*tasks)
-    open_ports = [(port, ports_services[port]) for port, is_open in zip(ports, results) if is_open]
+    open_ports = [(port, ports_services.get(port, 'Unknown'), banner) for port, (is_open, banner) in zip(ports, results) if is_open]
     return ip_str, True, open_ports
 
 async def full_scan(ip, ports_services):
@@ -54,10 +57,11 @@ async def full_scan(ip, ports_services):
     with tqdm(total=len(all_ports), desc=f"Full scan {ip}") as pbar:
         async def check(port):
             async with sem:
-                if await check_port(ip, port, timeout=0.5):
+                is_open, banner = await check_port(ip, port, timeout=0.5)
+                if is_open:
                     service = ports_services.get(port, 'Unknown')
                     pbar.update(1)
-                    return (port, service)
+                    return (port, service, banner)
                 else:
                     pbar.update(1)
                     return None
@@ -66,6 +70,18 @@ async def full_scan(ip, ports_services):
         open_ports = sorted([p for p in results if p is not None])
         return ip, open_ports
 
+async def get_dns_info(ip_str):
+    loop = asyncio.get_event_loop()
+    for attempt in range(3):
+        try:
+            hostname = (await loop.run_in_executor(None, socket.gethostbyaddr, ip_str))[0]
+            a_record = await loop.run_in_executor(None, socket.gethostbyname, hostname)
+            return hostname, a_record
+        except Exception as e:
+            print(f"Attempt {attempt+1} failed for {ip_str}: {e}")
+            await asyncio.sleep(2)
+    return "Unknown", "Unknown"
+
 async def main():
     local_ip = get_local_ip()
     network = ipaddress.IPv4Network(f"{local_ip}/24", strict=False)
@@ -161,6 +177,7 @@ async def main():
         1741: 'CiscoWorks',
         1812: 'RADIUS',
         1813: 'RADIUS Acct',
+        1984: 'Big Brother Monitor',
         1985: 'HSRP',
         2000: 'Cisco SCCP',
         2049: 'NFS',
@@ -172,6 +189,9 @@ async def main():
         2096: 'cPanel Webmail SSL',
         2100: 'Oracle XDB FTP',
         2222: 'DirectAdmin',
+        2242: 'Folio Remote Server',
+        2248: 'User Management Service',
+        2280: 'LNVPOLLER',
         2301: 'Compaq Insight',
         2381: 'HP Insight',
         2401: 'CVS',
@@ -180,6 +200,9 @@ async def main():
         2484: 'Oracle DB SSL',
         2638: 'Sybase',
         2710: 'XBT Tracker',
+        3000: 'Ruby on Rails',
+        3050: 'Interbase/Firebird',
+        3100: 'OpCon/xps',
         3128: 'Squid',
         3260: 'iSCSI',
         3268: 'AD Global Catalog',
@@ -192,9 +215,13 @@ async def main():
         3785: 'Ventrilo',
         3899: 'Remote Web Workplace',
         4000: 'Diablo II',
+        4190: 'Sieve Mail Filter',
         4321: 'BoBo',
+        4369: 'Erlang EPMD',
+        4431: 'adWISE Pipe',
         4664: 'Google Desktop',
         4899: 'Radmin',
+        4949: 'Munin',
         5000: 'UPnP',
         5001: 'Synology',
         5009: 'Airport Admin',
@@ -204,8 +231,10 @@ async def main():
         5222: 'XMPP',
         5269: 'XMPP SSL',
         5432: 'PostgreSQL',
+        54744: 'Unknown',
         5500: 'VNC Hotline',
         5554: 'Sasser',
+        5555: 'ADB',
         5631: 'pcAnywhere Data',
         5632: 'pcAnywhere Master',
         5800: 'VNC HTTP',
@@ -224,26 +253,38 @@ async def main():
         6697: 'IRC SSL',
         6881: 'BitTorrent',
         6969: 'BitTorrent',
+        7071: 'IWC Replica',
         7212: 'GhostSite',
         7777: 'Oracle Web',
         8000: 'HTTP Alternate',
+        8001: 'VCOM Tunnel',
+        8006: 'Proxmox VE',
         8008: 'HTTP Alternate',
+        8009: 'AJP',
         8080: 'HTTP Alternate',
         8081: 'HTTP Alternate',
+        8088: 'HTTP Alternate',
+        8089: 'Splunk',
+        8090: 'HTTP Alternate',
         8118: 'Privoxy',
         8200: 'MiniDLNA',
+        8424: 'Unknown',
         8443: 'HTTPS Alternate',
         8554: 'RTSP Alternate',
+        8762: 'Unknown',
         8767: 'TeamSpeak',
         8888: 'HTTP Alternate',
         9000: 'SonarQube',
         9001: 'Tor',
+        9002: 'Dynamix',
+        9003: 'Unknown',
         9030: 'Tor',
         9040: 'Tor',
         9050: 'Tor',
         9100: 'JetDirect',
         9119: 'MXit',
         9200: 'Elasticsearch',
+        9230: 'Unknown',
         9418: 'Git',
         9999: 'Abyss',
         10000: 'Webmin',
@@ -289,19 +330,31 @@ async def main():
     results = await asyncio.gather(*tasks)
 
     active_devices = []
-    devices_need_full = []
     common_open_ports = {}
     for ip, active, open_ports in results:
         if active:
             active_devices.append(ip)
             if open_ports:
-                devices_need_full.append(ip)
-                common_open_ports[ip] = open_ports
+                common_open_ports[str(ip)] = open_ports
+
+    dns_tasks = [get_dns_info(str(ip)) for ip in active_devices]
+    dns_results = await asyncio.gather(*dns_tasks)
+    dns_infos = {str(active_devices[i]): dns_results[i] for i in range(len(active_devices))}
+
+    devices_need_full = []
+    for ip in active_devices:
+        ip_str = str(ip)
+        if common_open_ports.get(ip_str) or dns_infos[ip_str][0] != "Unknown":
+            devices_need_full.append(ip_str)
 
     full_tasks = [full_scan(ip, ports_services) for ip in devices_need_full]
     full_results = await asyncio.gather(*full_tasks)
 
-    devices_with_open_ports = {ip: common_open_ports.get(ip, []) for ip in active_devices}
+    devices_with_open_ports = {}
+    for ip in active_devices:
+        ip_str = str(ip)
+        devices_with_open_ports[ip_str] = common_open_ports.get(ip_str, [])
+
     for full_ip, full_open in full_results:
         devices_with_open_ports[full_ip] = full_open
 
@@ -311,9 +364,13 @@ async def main():
 
     print("\nDevices with open ports:")
     for ip, open_ports in devices_with_open_ports.items():
-        if open_ports:
-            print(f"{ip}:")
-            for port, service in open_ports:
-                print(f"  - {port}: {service}")
+        hostname, a_record = dns_infos[ip]
+        print(f"{ip} (Hostname: {hostname}, A: {a_record}):")
+        if hostname != "Unknown":
+            print(f"  DNS Records:")
+            print(f"    PTR: {hostname}")
+            print(f"    A: {a_record}")
+        for port, service, banner in open_ports:
+            print(f"  - {port}: {service}" + (f" ({banner})" if banner else ""))
 
 asyncio.run(main())