# zschema sub-schema for zgrab2's smb module
# Registers zgrab2-smb globally, and smb with the main zgrab2 schema.
from zschema.leaves import *
from zschema.compounds import *
import zschema.registry

import zcrypto_schemas.zcrypto as zcrypto
from . import zgrab2

header_log = {
    "protocol_id": Binary(),
    "status": Unsigned32BitInteger(),
    "command": Unsigned16BitInteger(),
    "credits": Unsigned16BitInteger(),
    "flags": Unsigned32BitInteger(),
}


# Return a (shallow) copy of base, with the fields of new merged atop it
def extended(base, new):
    copy = {k: v for k, v in base.items()}
    for k, v in new.items():
        copy[k] = v
    return copy


negotiate_log = SubRecord(
    extended(
        header_log,
        {
            "security_mode": Unsigned16BitInteger(),
            "dialect_revision": Unsigned16BitInteger(),
            "server_guid": Binary(),
            "capabilities": Unsigned32BitInteger(),
            "system_time": Unsigned32BitInteger(),
            "server_start_time": Unsigned32BitInteger(),
            "authentication_types": ListOf(String()),
        },
    )
)

session_setup_log = SubRecord(
    extended(
        header_log,
        {
            "setup_flags": Unsigned16BitInteger(),
            "target_name": String(),
            "negotiate_flags": Unsigned32BitInteger(),
        },
    )
)


smb_scan_response = SubRecord(
    {
        "result": SubRecord(
            {
                "smbv1_support": Boolean(),
                "smb_version": SubRecord(
                    {
                        "major": Unsigned8BitInteger(doc="Major version"),
                        "minor": Unsigned8BitInteger(doc="Minor version"),
                        "revision": Unsigned8BitInteger(doc="Protocol Revision"),
                        "version_string": String(doc="Full SMB Version String"),
                    }
                ),
                "native_os": String(doc="Operating system claimed by server"),
                "ntlm": String(doc="Native LAN Manager"),
                "group_name": String(doc="Group name"),
                "smb_capabilities": SubRecord(
                    {
                        "smb_dfs_support": Boolean(
                            doc="Server supports Distributed File System"
                        ),
                        "smb_leasing_support": Boolean(doc="Server supports Leasing"),
                        "smb_multicredit_support": Boolean(
                            doc="Server supports multi-credit operations"
                        ),
                        "smb_multchan_support": Boolean(
                            doc="Server supports multiple channels per session"
                        ),
                        "smb_persistent_handle_support": Boolean(
                            doc="Server supports persistent handles"
                        ),
                        "smb_directory_leasing_support": Boolean(
                            doc="Server supports directory leasing"
                        ),
                        "smb_encryption_support": Boolean(
                            doc="Server supports encryption"
                        ),
                    },
                    doc="Capabilities flags for the connection. See [MS-SMB2] Sect. 2.2.4.",
                ),
                "negotiation_log": negotiate_log,
                "has_ntlm": Boolean(
                    doc="Server supports the NTLM authentication method"
                ),
                "session_setup_log": session_setup_log,
            }
        )
    },
    extends=zgrab2.base_scan_response,
)

zschema.registry.register_schema("zgrab2-smb", smb_scan_response)

zgrab2.register_scan_response_type("smb", smb_scan_response)