from zschema.leaves import *
from zschema.compounds import *
import zschema.registry
from collections import defaultdict

# Base / shared schema types for zgrab2

import zcrypto_schemas.zcrypto as zcrypto

# Map of protocol-name -> protocl-schema. This is wrapped in a SubRecord
# Protocols are responsible for calling register_scan_response_type(protocol_name, schema).
# Failure to do so will result in a validation exception for any scan results containing that protocol.
# NOTE: Scans with custom names will cause the validator to fail.
# TODO: Can this be somehow replaced with the main registry?
scan_response_types = {}


# Placeholder / RFU. Many mysql fields are only included in debug mode.
def DebugOnly(childType):
    return childType


# Get a dict of all keys, mapping key -> true
def FlagsSet(keys, **kwargs):
    return SubRecord({key: Boolean() for key in keys}, **kwargs)


# zgrab2/processing.go: Grab
grab_result = Record(
    {
        # TODO: ip may be required; see https://github.com/zmap/zgrab2/issues/104
        "ip": IPv4Address(required=False, doc="The IP address of the target."),
        "domain": String(
            required=False, doc="The domain name of the target, if available."
        ),
        "data": SubRecord(scan_response_types, doc="The scan data for this host."),
    }
)

# zgrab2/module.go: const SCAN_*
STATUS_VALUES = [
    "success",
    "connection-refused",
    "connection-timeout",
    "connection-closed",
    "io-timeout",
    "protocol-error",
    "application-error",
    "unknown-error",
]

# zgrab2/module.go: ScanResponse
base_scan_response = SubRecord(
    {
        "status": Enum(values=STATUS_VALUES, doc="The status of the request."),
        "protocol": String(doc="The identifier of the protocol being scanned."),
        "timestamp": DateTime(doc="The time the scan was started."),
        "result": SubRecord(
            {}, required=False
        ),  # This is overridden by the protocols' implementations
        "error": String(
            required=False,
            doc="If the status was not success, error may contain information about the failure.",
        ),
        # TODO: error_component? domain?
    }
)

# zgrab2/tls.go: TLSLog
tls_log = SubRecord(
    {
        "handshake_log": zcrypto.TLSHandshake(doc="The TLS handshake log."),
        "heartbleed_log": zcrypto.HeartbleedLog(
            doc="The heartbleed scan log, if heartbleed scanning was enabled; otherwise, absent."
        ),
    }
)


# Register a schema type for responses with the given name.
def register_scan_response_type(name, schema):
    scan_response_types[name] = schema


zschema.registry.register_schema("zgrab2", grab_result)