"""Synthetic protocol frames only: never connect to remote scan targets."""
import socket
import struct
import threading
import unittest
from unittest.mock import MagicMock, patch

from rpcscope.smb_negotiate import negotiate_request, parse_negotiate, probe_smb_compression


def context(kind, data):
    value = struct.pack("<HHI", kind, len(data), 0) + data
    return value + b"\x00" * (-len(value) % 8)


def response(dialect=0x311, algorithms=(1,), compression=True):
    data = bytearray(128)
    data[:4] = b"\xfeSMB"
    struct.pack_into("<H", data, 4, 64)
    struct.pack_into("<I", data, 16, 1)
    struct.pack_into("<HHHH", data, 64, 65, 1, dialect, 1 + compression if dialect == 0x311 else 0)
    if dialect == 0x311:
        struct.pack_into("<I", data, 124, 128)
        data += context(1, struct.pack("<HHH", 1, 0, 1))
        if compression:
            data += context(3, struct.pack("<HHI", len(algorithms), 0, 0) + struct.pack(f"<{len(algorithms)}H", *algorithms))
    return bytes(data)


class SmbNegotiateTests(unittest.TestCase):
    def test_request_is_small_normal_negotiate_with_correct_context_offsets(self):
        frame = negotiate_request()
        size = int.from_bytes(frame[:4], "big")
        self.assertEqual(size, len(frame) - 4)
        self.assertLess(size, 256)
        data = frame[4:]
        self.assertEqual(data[:4], b"\xfeSMB")
        self.assertEqual(struct.unpack_from("<H", data, 12)[0], 0)
        self.assertEqual(struct.unpack_from("<H", data, 64)[0], 36)
        offset, count = struct.unpack_from("<IH", data, 92)
        self.assertEqual((offset, count), (112, 2))
        self.assertEqual(struct.unpack_from("<HHI", data, offset), (1, 38, 0))
        self.assertEqual(struct.unpack_from("<HHI", data, 160), (3, 14, 0))

    def test_valid_capability_is_exposure_not_cve(self):
        for algorithms in ((1,), (1, 2, 3), (0, 1)):
            with self.subTest(algorithms=algorithms):
                outcome = parse_negotiate(response(algorithms=algorithms))
                self.assertEqual(outcome.status, "exposed")
                self.assertIn("patch state unknown", outcome.evidence)

    def test_absent_capability_is_not_a_patch_verdict(self):
        for frame in (response(0x202), response(0x210), response(0x300), response(0x302),
                      response(compression=False), response(algorithms=(0,))):
            outcome = parse_negotiate(frame)
            self.assertEqual(outcome.status, "not-exposed")
            self.assertIn("not a patch/CVE verdict", outcome.evidence)

    def test_malformed_offsets_lengths_and_wrong_responses_never_become_negative(self):
        modifications = [(0, b"\xffSMB"), (4, b"\x20\x00"), (8, b"\x22\x00\x00\xc0"),
                         (12, b"\x01\x00"), (16, b"\x00\x00\x00\x00"),
                         (16, b"\x03\x00\x00\x00"), (20, b"\x80\x00\x00\x00"),
                         (24, b"\x01"), (64, b"\x40\x00"), (68, b"\xff\xff"),
                         (70, b"\xff\xff"), (124, b"\x01\x00\x00\x00"),
                         (124, b"\x80\xff\x00\x00"), (130, b"\xff\xff"),
                         (136, b"\x02"), (140, b"\x02"),
                         (144, b"\x01\x00"), (152, b"\xff\xff")]
        for offset, value in modifications:
            with self.subTest(offset=offset, value=value):
                data = bytearray(response())
                data[offset:offset + len(value)] = value
                with self.assertRaises(ValueError):
                    parse_negotiate(bytes(data))
        for end in (0, 4, 64, 127, 140, 155):
            with self.subTest(end=end), self.assertRaises(ValueError):
                parse_negotiate(response()[:end])

    def test_older_dialect_reserved_fields_are_ignored(self):
        for dialect in (0x202, 0x210, 0x300, 0x302):
            with self.subTest(dialect=dialect):
                data = bytearray(response(dialect))
                struct.pack_into("<H", data, 70, 65535)
                struct.pack_into("<I", data, 124, 0xffffffff)
                self.assertEqual(parse_negotiate(bytes(data)).status, "not-exposed")

    def test_padding_reserved_bytes_and_chained_support_do_not_hide_capability(self):
        data = bytearray(response())
        struct.pack_into("<I", data, 132, 0xffffffff)
        struct.pack_into("<H", data, 154, 0xffff)
        struct.pack_into("<I", data, 156, 1)
        self.assertEqual(parse_negotiate(bytes(data)).status, "exposed")

    def test_server_algorithm_list_uses_intersection_with_our_offer(self):
        self.assertEqual(parse_negotiate(response(algorithms=(1, 4, 5))).status, "exposed")
        self.assertEqual(parse_negotiate(response(algorithms=(4, 5))).status, "not-exposed")
        with self.assertRaises(ValueError):
            parse_negotiate(response(algorithms=(32,)))

    def test_duplicate_unknown_contexts_are_ignored_but_known_duplicates_are_invalid(self):
        data = bytearray(response())
        struct.pack_into("<H", data, 70, 4)
        data += context(0x100, b"") + context(0x100, b"")
        self.assertEqual(parse_negotiate(bytes(data)).status, "exposed")
        data = bytearray(response())
        struct.pack_into("<H", data, 70, 3)
        data += context(3, struct.pack("<HHIH", 1, 0, 0, 1))
        with self.assertRaises(ValueError):
            parse_negotiate(bytes(data))

    def test_fragmented_tcp_response_is_read_completely_without_login(self):
        payload = response()
        frame = len(payload).to_bytes(4, "big") + payload
        connection = MagicMock()
        connection.recv.side_effect = [bytes([byte]) for byte in frame]
        with patch("rpcscope.smb_negotiate.socket.create_connection") as connect:
            connect.return_value.__enter__.return_value = connection
            outcome = probe_smb_compression("192.0.2.1")
        self.assertEqual(outcome.status, "exposed")
        self.assertEqual(connection.sendall.call_count, 1)
        self.assertEqual(connection.sendall.call_args.args[0][4:8], b"\xfeSMB")
        self.assertEqual(connect.call_args.args, (("192.0.2.1", 445),))

    def test_cancellation_does_not_connect(self):
        cancel = threading.Event()
        cancel.set()
        with patch("rpcscope.smb_negotiate.socket.create_connection") as connect:
            self.assertEqual(probe_smb_compression("192.0.2.1", cancel=cancel).status, "skipped-cancelled")
            connect.assert_not_called()

    def test_timeout_eof_or_large_frame_remain_inconclusive(self):
        with patch("rpcscope.smb_negotiate.socket.create_connection", side_effect=socket.timeout("timed out")):
            self.assertEqual(probe_smb_compression("192.0.2.1").status, "inconclusive")
        for value in (b"", b"\x01\x00\x00\x00", b"\x00\x01\x00\x01"):
            connection = MagicMock()
            connection.recv.return_value = value
            with patch("rpcscope.smb_negotiate.socket.create_connection") as connect:
                connect.return_value.__enter__.return_value = connection
                self.assertEqual(probe_smb_compression("192.0.2.1").status, "inconclusive")
                self.assertEqual(connection.recv.call_count, 1)

    def test_total_deadline_not_reset_per_fragment(self):
        connection = MagicMock()
        with patch("rpcscope.smb_negotiate.socket.create_connection") as connect, \
                patch("rpcscope.smb_negotiate.time.monotonic", side_effect=[0, 0, 0, 10]):
            connect.return_value.__enter__.return_value = connection
            outcome = probe_smb_compression("192.0.2.1", timeout=5)
        self.assertEqual(outcome.status, "inconclusive")
        self.assertIn("timeout", outcome.evidence)
        connection.recv.assert_not_called()