[DEV] Database model revamp

Merge pull request #15 from MatrixEditor/dev/db-revamp

Author: MatrixEditor

dementor/__init__.py

@@ -18,5 +18,5 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-__version__ = "1.0.0.dev13"
+__version__ = "1.0.0.dev14"
 __author__ = "MatrixEditor"

dementor/assets/Dementor.toml

@@ -186,10 +186,10 @@ UPnP = true
 #   These extra fields (e.g., `foo`, `bar`) will be passed to the final
 #   FilterObj instance and may be used for protocol-specific handling.
 
-# AnswerTo = []
+# Target = []
 
 # Describes a list of hosts to *exclude* from poisoning (blacklist approach).
-# Uses the same filter structure as `AnswerTo`.
+# Uses the same filter structure as Target.
 
 # Ignore = []
 
@@ -205,16 +205,10 @@ UPnP = true
 
 DuplicateCreds = true
 
-# Custom database filename (overrides default "Dementor.db"). Do not put any
-# path information here.
-
-# Name = "database.db"
-
-# Custom path where the database file will be stored. If set, it overrides the
-# workspace directory.
-
-# Directory = "/path/to/directory/"
-
+# Dialect = "sqlite"
+# Driver = "pysqlite"
+# Url = "sqlite:///:memory:"
+# Path = "Dementor.db"
 
 # =============================================================================
 # mDNS

dementor/config/session.py

@@ -17,17 +17,43 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
+# pyright: reportUninitializedInstanceVariable=false
 import asyncio
+import typing
 
-from typing import Any, List
-
+from typing import Any
 from pathlib import Path
 
 from dementor.config.toml import TomlConfig, Attribute
 from dementor.config.util import is_true
 from dementor.paths import DEMENTOR_PATH
 from dementor import config
 
+if typing.TYPE_CHECKING:
+    from dementor.protocols import (
+        kerberos,
+        mdns,
+        netbios,
+        llmnr,
+        ldap,
+        smb,
+        smtp,
+        ftp,
+        http,
+        imap,
+        ipp,
+        mssql,
+        mysql,
+        pop3,
+        quic,
+        ssdp,
+        upnp,
+        x11,
+    )
+    from dementor.protocols.msrpc import rpc
+    from dementor.db.model import DementorDB
+    from dementor.db.connector import DatabaseConfig
+
 
 class SessionConfig(TomlConfig):
     _section_ = "Dementor"
@@ -64,14 +90,66 @@ class SessionConfig(TomlConfig):
     ]
 
     # TODO: move into .pyi
-    db: Any
-    db_config: Any
-    krb5_config: Any
-    mdns_config: Any
-    llmnr_config: Any
-    quic_config: Any
-    netbiosns_config: Any
-    ldap_config: List[Any]
+    if typing.TYPE_CHECKING:
+        workspace_path: str
+        extra_modules: list[str]
+        ipv6: str | None
+        ipv4: str | None
+        interface: str | None
+        protocols: dict[str, Any]
+
+        db: DementorDB
+        db_config: DatabaseConfig
+        krb5_config: kerberos.KerberosConfig
+        mdns_config: mdns.MDNSConfig
+        llmnr_config: llmnr.LLMNRConfig
+        netbiosns_config: netbios.NBTNSConfig
+        ldap_config: list[ldap.LDAPServerConfig]
+        smtp_servers: list[smtp.SMTPServerConfig]
+        smb_config: list[smb.SMBServerConfig]
+        ftp_config: list[ftp.FTPServerConfig]
+        proxy_config: http.ProxyAutoConfig
+        http_config: list[http.HTTPServerConfig]
+        winrm_config: list[http.HTTPServerConfig]
+        imap_config: list[imap.IMAPServerConfig]
+        ipp_config: ipp.IPPConfig
+        rpc_config: rpc.RPCConfig
+        mssql_config: mssql.MSSQLConfig
+        ssrp_config: mssql.SSRPConfig
+        mysql_config: mysql.MySQLConfig
+        pop3_config: list[pop3.POP3ServerConfig]
+        quic_config: quic.QuicServerConfig
+        ssdp_config: ssdp.SSDPConfig
+        upnp_config: upnp.UPNPConfig
+        x11_config: x11.X11Config
+
+        ntlm_challange: bytes
+        ntlm_ess: bool
+        analysis: bool
+        loop: asyncio.AbstractEventLoop
+
+        llmnr_enabled: bool
+        nbtns_enabled: bool
+        nbtds_enabled: bool
+        smtp_enabled: bool
+        smb_enabled: bool
+        ftp_enabled: bool
+        kdc_enabled: bool
+        ldap_enabled: bool
+        quic_enabled: bool
+        mdns_enabled: bool
+        http_enabled: bool
+        rpc_enabled: bool
+        winrm_enabled: bool
+        mssql_enabled: bool
+        ssrp_enabled: bool
+        imap_enabled: bool
+        pop3_enabled: bool
+        mysql_enabled: bool
+        x11_enabled: bool
+        ipp_enabled: bool
+        ssdp_enabled: bool
+        upnp_enabled: bool
 
     def __init__(self) -> None:
         super().__init__(config._get_global_config().get("Dementor", {}))
@@ -90,9 +168,6 @@ def __init__(self) -> None:
         self.ntlm_challange = b"1337LEET"
         self.ntlm_ess = True
 
-        # SMB configuration
-        self.smb_server_config = []
-
     def is_bound_to_all(self) -> bool:
         # REVISIT: this should raise an exception
         return self.interface == "ALL"

dementor/config/toml.py

@@ -17,11 +17,12 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
-from typing import NamedTuple, Callable, Any
+from typing import NamedTuple, Callable, Any, TypeVar
 
 from dementor.config.util import get_value
 
 _LOCAL = object()
+_T = TypeVar("_T", bound="TomlConfig")
 
 
 class Attribute(NamedTuple):
@@ -36,7 +37,7 @@ class TomlConfig:
     _section_: str | None
     _fields_: list[Attribute]
 
-    def __init__(self, config: dict | None = None) -> None:
+    def __init__(self, config: dict[str, Any] | None = None) -> None:
         for field in self._fields_:
             self._set_field(
                 config or {},
@@ -62,8 +63,12 @@ def __getitem__(self, key: str) -> Any:
         raise KeyError(f"Could not find config with key {key!r}")
 
     @staticmethod
-    def build_config(cls_ty, section: str | None = None) -> Any:
-        return cls_ty(get_value(section or cls_ty._section_, key=None, default={}))
+    def build_config(cls_ty: type[_T], section: str | None = None) -> _T:
+        section_name = section or cls_ty._section_
+        if section_name is None:
+            raise ValueError("section cannot be None")
+
+        return cls_ty(get_value(section_name, key=None, default={}))
 
     def _set_field(
         self,

dementor/config/util.py

@@ -53,10 +53,10 @@ def is_true(value: str) -> bool:
 
 
 class BytesValue:
-    def __init__(self, length=None) -> None:
-        self.length = length
+    def __init__(self, length: int | None = None) -> None:
+        self.length: int | None = length
 
-    def __call__(self, value) -> Any:
+    def __call__(self, value: Any) -> bytes:
         match value:
             case None:
                 return secrets.token_bytes(self.length or 1)
@@ -78,7 +78,7 @@ def random_value(size: int) -> str:
     return "".join(random.choice(string.ascii_letters) for _ in range(size))
 
 
-def format_string(value: str, locals: dict | None = None) -> str:
+def format_string(value: str, locals: dict[str, Any] | None = None) -> str:
     config = _get_global_config()
     try:
         template = _SANDBOX.from_string(value)

dementor/db/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2025-Present MatrixEditor
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+_CLEARTEXT = "Cleartext"
+_NO_USER = "<missing-user>"
+
+
+def normalize_client_address(client: str) -> str:
+    return client.removeprefix("::ffff:")
+

dementor/db/connector.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2025-Present MatrixEditor
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import sqlite3
+from sqlalchemy import Engine, create_engine
+
+from dementor.config.session import SessionConfig
+from dementor.db.model import DementorDB, ModelBase
+from dementor.log.logger import dm_logger
+from dementor.config.toml import TomlConfig, Attribute as A
+
+
+class DatabaseConfig(TomlConfig):
+    _section_ = "DB"
+    _fields_ = [
+        A("db_raw_path", "Url", None),
+        A("db_path", "Path", "Dementor.db"),
+        A("db_duplicate_creds", "DuplicateCreds", False),
+        A("db_dialect", "Dialect", None),
+        A("db_driver", "Driver", None),
+    ]
+
+
+def init_dementor_db(session: SessionConfig) -> Engine | None:
+    engine = init_engine(session)
+    if engine is not None:
+        ModelBase.metadata.create_all(engine)
+    return engine
+
+
+def init_engine(session: SessionConfig) -> Engine | None:
+    # based on dialect and driver configuration
+    raw_path = session.db_config.db_raw_path
+    if raw_path is None:
+        # fall back to constructing the path manually
+        dialect = session.db_config.db_dialect or "sqlite"
+        driver = session.db_config.db_driver or "pysqlite"
+        path = session.db_config.db_path
+        if not path:
+            return dm_logger.error("Database path not specified!")
+
+        if dialect == "sqlite":
+            if path != ":memory:":
+                path = f"/{session.resolve_path(path)}"
+
+        # see https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls
+        raw_path = f"{dialect}+{driver}://{path}"
+    else:
+        sql_type, path = raw_path.split("://")
+        if sql_type.count("+") > 0:
+            dialect, driver = sql_type.split("+")
+        else:
+            dialect = sql_type
+            driver = "<default>"
+
+    if dialect != "sqlite":
+        first_element, *parts = path.split("/")
+        if "@" in first_element:
+            first_element = first_element.split("@")[1]
+            path = "***:***@" + "/".join([first_element] + list(parts))
+
+    dm_logger.debug("Using database [%s:%s] at: %s", dialect, driver, path)
+    return create_engine(raw_path, isolation_level="AUTOCOMMIT", future=True)
+
+
+def create_db(session: SessionConfig) -> DementorDB:
+    # TODO: add support for custom database implementations
+    engine = init_engine(session)
+    if not engine:
+        raise Exception("Failed to create database engine")
+    return DementorDB(engine, session)