# zschema sub-schema for zgrab2's postgres module
# Registers zgrab2-postgres globally, and postgres 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

# modules/postgres/scanner.go - decodeError() (TODO: Currently an unconstrained
# map[string]string; it is possible to get "unknown (0x%x)" fields, but it
# would probably be proper to reject those at this point)

# These are defined in detail at
#   https://www.postgresql.org/docs/10/static/protocol-error-fields.html
postgres_error = SubRecord(
    {
        "severity": WhitespaceAnalyzedString(),
        "severity_v": WhitespaceAnalyzedString(),
        "code": WhitespaceAnalyzedString(),
        "message": WhitespaceAnalyzedString(),
        "detail": WhitespaceAnalyzedString(),
        "hint": WhitespaceAnalyzedString(),
        "position": WhitespaceAnalyzedString(),
        "internal_position": WhitespaceAnalyzedString(),
        "internal_query": WhitespaceAnalyzedString(),
        "where": WhitespaceAnalyzedString(),
        "schema": WhitespaceAnalyzedString(),
        "table": WhitespaceAnalyzedString(),
        "data": WhitespaceAnalyzedString(),
        "constraint": WhitespaceAnalyzedString(),
        "file": WhitespaceAnalyzedString(),
        "line": WhitespaceAnalyzedString(),
        "routine": WhitespaceAnalyzedString(),
        "_unknown_error_tag": WhitespaceAnalyzedString(),
    }
)

# modules/postgres/scanner.go - decodeAuthMode()
AUTH_MODES = [
    "kerberos_v5",
    "password_cleartext",
    "password_md5",
    "scm_credentials",
    "gss",
    "sspi",
    "sasl",
    "ok",
    "gss-continue",
    "sasl-continue",
    "sasl-final",
]

# modules/postgres/scanner.go: AuthenticationMode
postgres_auth_mode = SubRecord(
    {
        "mode": Enum(values=AUTH_MODES, required=False),  # this gets lifted
        "Payload": Binary(),
    }
)

# modules/postgres/scanner.go: BackendKeyData
postgres_key_data = SubRecord(
    {
        "process_id": Unsigned32BitInteger(),
        "secret_key": Unsigned32BitInteger(),
    }
)

# modules/postgres/scanner.go: PostgresResults
postgres_scan_response = SubRecord(
    {
        "result": SubRecord(
            {
                "tls": zgrab2.tls_log,
                "supported_versions": WhitespaceAnalyzedString(),
                "protocol_error": postgres_error,
                "startup_error": postgres_error,
                "is_ssl": Boolean(),
                "authentication_mode": postgres_auth_mode,
                # TODO FIXME: This is currendly an unconstrained map[string]string
                "server_parameters": WhitespaceAnalyzedString(),
                "backend_key_data": postgres_key_data,
                "transaction_status": WhitespaceAnalyzedString(),
            }
        )
    },
    extends=zgrab2.base_scan_response,
)

zschema.registry.register_schema("zgrab2-postgres", postgres_scan_response)

zgrab2.register_scan_response_type("postgres", postgres_scan_response)