# See the file "COPYING" in the main distribution directory for copyright.

module LDAP;

import ASN1;
import spicy;

# https://tools.ietf.org/html/rfc4511#
# https://ldap.com/ldapv3-wire-protocol-reference-asn1-ber/
# https://lapo.it/asn1js

#- Operation opcode ----------------------------------------------------------
public type ProtocolOpcode = enum {
  BIND_REQUEST            = 0,
  BIND_RESPONSE           = 1,
  UNBIND_REQUEST          = 2,
  SEARCH_REQUEST          = 3,
  SEARCH_RESULT_ENTRY     = 4,
  SEARCH_RESULT_DONE      = 5,
  MODIFY_REQUEST          = 6,
  MODIFY_RESPONSE         = 7,
  ADD_REQUEST             = 8,
  ADD_RESPONSE            = 9,
  DEL_REQUEST             = 10,
  DEL_RESPONSE            = 11,
  MOD_DN_REQUEST          = 12,
  MOD_DN_RESPONSE         = 13,
  COMPARE_REQUEST         = 14,
  COMPARE_RESPONSE        = 15,
  ABANDON_REQUEST         = 16,
  SEARCH_RESULT_REFERENCE = 19,
  EXTENDED_REQUEST        = 23,
  EXTENDED_RESPONSE       = 24,
  INTERMEDIATE_RESPONSE   = 25,
};

#- Result code ---------------------------------------------------------------
public type ResultCode = enum {
  SUCCESS                         = 0,
  OPERATIONS_ERROR                = 1,
  PROTOCOL_ERROR                  = 2,
  TIME_LIMIT_EXCEEDED             = 3,
  SIZE_LIMIT_EXCEEDED             = 4,
  COMPARE_FALSE                   = 5,
  COMPARE_TRUE                    = 6,
  AUTH_METHOD_NOT_SUPPORTED       = 7,
  STRONGER_AUTH_REQUIRED          = 8,
  PARTIAL_RESULTS                 = 9,
  REFERRAL                        = 10,
  ADMIN_LIMIT_EXCEEDED            = 11,
  UNAVAILABLE_CRITICAL_EXTENSION  = 12,
  CONFIDENTIALITY_REQUIRED        = 13,
  SASL_BIND_IN_PROGRESS           = 14,
  NO_SUCH_ATTRIBUTE               = 16,
  UNDEFINED_ATTRIBUTE_TYPE        = 17,
  INAPPROPRIATE_MATCHING          = 18,
  CONSTRAINT_VIOLATION            = 19,
  ATTRIBUTE_OR_VALUE_EXISTS       = 20,
  INVALID_ATTRIBUTE_SYNTAX        = 21,
  NO_SUCH_OBJECT                  = 32,
  ALIAS_PROBLEM                   = 33,
  INVALID_DNSYNTAX                = 34,
  ALIAS_DEREFERENCING_PROBLEM     = 36,
  INAPPROPRIATE_AUTHENTICATION    = 48,
  INVALID_CREDENTIALS             = 49,
  INSUFFICIENT_ACCESS_RIGHTS      = 50,
  BUSY                            = 51,
  UNAVAILABLE                     = 52,
  UNWILLING_TO_PERFORM            = 53,
  LOOP_DETECT                     = 54,
  SORT_CONTROL_MISSING            = 60,
  OFFSET_RANGE_ERROR              = 61,
  NAMING_VIOLATION                = 64,
  OBJECT_CLASS_VIOLATION          = 65,
  NOT_ALLOWED_ON_NON_LEAF         = 66,
  NOT_ALLOWED_ON_RDN              = 67,
  ENTRY_ALREADY_EXISTS            = 68,
  OBJECT_CLASS_MODS_PROHIBITED    = 69,
  RESULTS_TOO_LARGE               = 70,
  AFFECTS_MULTIPLE_DSAS           = 71,
  CONTROL_ERROR                   = 76,
  OTHER                           = 80,
  SERVER_DOWN                     = 81,
  LOCAL_ERROR                     = 82,
  ENCODING_ERROR                  = 83,
  DECODING_ERROR                  = 84,
  TIMEOUT                         = 85,
  AUTH_UNKNOWN                    = 86,
  FILTER_ERROR                    = 87,
  USER_CANCELED                   = 88,
  PARAM_ERROR                     = 89,
  NO_MEMORY                       = 90,
  CONNECT_ERROR                   = 91,
  NOT_SUPPORTED                   = 92,
  CONTROL_NOT_FOUND               = 93,
  NO_RESULTS_RETURNED             = 94,
  MORE_RESULTS_TO_RETURN          = 95,
  CLIENT_LOOP                     = 96,
  REFERRAL_LIMIT_EXCEEDED         = 97,
  INVALID_RESPONSE                = 100,
  AMBIGUOUS_RESPONSE              = 101,
  TLS_NOT_SUPPORTED               = 112,
  INTERMEDIATE_RESPONSE           = 113,
  UNKNOWN_TYPE                    = 114,
  LCUP_INVALID_DATA               = 115,
  LCUP_UNSUPPORTED_SCHEME         = 116,
  LCUP_RELOAD_REQUIRED            = 117,
  CANCELED                        = 118,
  NO_SUCH_OPERATION               = 119,
  TOO_LATE                        = 120,
  CANNOT_CANCEL                   = 121,
  ASSERTION_FAILED                = 122,
  AUTHORIZATION_DENIED            = 123,
};

#-----------------------------------------------------------------------------
public type Result = unit {
  code: ASN1::ASN1Message(True) &convert=cast<ResultCode>(cast<uint8>($$.body.num_value))
                                &default=ResultCode::Undef;
  matchedDN: ASN1::ASN1Message(True) &convert=$$.body.str_value
                                     &default="";
  diagnosticMessage: ASN1::ASN1Message(True) &convert=$$.body.str_value
                                             &default="";

  # TODO: if we want to parse referral URIs in result
  # https://tools.ietf.org/html/rfc4511#section-4.1.10
};

# 1.2.840.48018.1.2.2 (MS KRB5 - Microsoft Kerberos 5)
const GSSAPI_MECH_MS_KRB5 = "1.2.840.48018.1.2.2";

# Supported SASL stripping modes.
type MessageMode = enum {
  MS_KRB5 = 1, # Payload starts with a 4 byte length followed by a wrap token that may or may not be sealed.
  TLS = 2,     # Client/server used StartTLS, forward to SSL analyzer.
  MAYBE_ENCRYPTED = 3,  # Use a heuristic to determine encrypted traffic.
  CLEARTEXT = 4,  # Assume cleartext.
  ENCRYPTED = 5,  # Assume encrypted.
};

type Ctx = struct {
  messageMode: MessageMode;   # Message dispatching mode
  saslMechanism: string;      # The SASL mechanism selected by the client.
  startTlsRequested: bool;    # Did the client use the StartTLS extended request?
};

#-----------------------------------------------------------------------------
public type Messages = unit {
  %context = Ctx;
  : MessageDispatch(self.context())[];
};

#-----------------------------------------------------------------------------
public type MessageDispatch = unit(ctx: Ctx&) {
  switch( ctx.messageMode ) {
    MessageMode::Undef -> : Message(ctx);
    MessageMode::MS_KRB5 -> : SaslMsKrb5Stripper(ctx);
    MessageMode::TLS -> : TlsForward;  # never returns
    MessageMode::MAYBE_ENCRYPTED -> : MaybeEncrypted(ctx);
    MessageMode::CLEARTEXT -> : Message(ctx);
    MessageMode::ENCRYPTED -> : EncryptedMessage;
  };
};


#-----------------------------------------------------------------------------
type MaybeEncrypted = unit(ctx: Ctx&) {
  # A plaintext LDAP message always starts with at least 3 bytes and the first
  # byte is 0x30 for the sequence. A SASL encrypted message starts with a 4 byte
  # length field. The heuristic here is that if the first byte is a 0x30,
  # assume it's unencrypted LDAP. This should be pretty good, if it was an
  # encrypted/SASL wrapped message, it would have a size between 0x30000000 and
  # 0x30FFFFFF, meaning at least a size of ~768MB, which seems unlikely.
  var start: iterator<stream>;
  var saslLen: uint64;
  var mech: bytes;

  on %init {
    self.start = self.input();
    # Don't have starts_with() on string, work around that.
    # https://github.com/zeek/spicy/issues/1807
    self.mech = ctx.saslMechanism.encode(spicy::Charset::UTF8);
  }

  first: uint8 {
    if ( $$ == 0x30 ) {
      ctx.messageMode = MessageMode::CLEARTEXT;
    } else {
      ctx.messageMode = MessageMode::ENCRYPTED;
    }
  }

  # As a further heuristic, if encrypted mode was decided and the client
  # requested GSSAPI or GSS-SPNEGO (or we just didn't see it) peak a bit
  # into the SASL payload and check if it starts with a 0504 (WRAP_TOKEN).
  # If so, switch into KRB mode assuming that's what is being used and
  # have a chance seeing some more plaintext LDAP in non-sealed tokens.
  rem: uint8[3] if ( ctx.messageMode == MessageMode::ENCRYPTED && (|self.mech| == 0 || self.mech.starts_with(b"GSS")) ) {
    self.saslLen = (uint64(self.first) << 24) + (uint64($$[0]) << 16) + (uint64($$[1]) << 8) + uint64($$[2]);
  }

  : uint16 if ( self.saslLen >= 2 ) {
    if ( $$ == 0x0504 ) {
      ctx.messageMode = MessageMode::MS_KRB5;
    }
  }

  # Rewind the input.
  : void {
    # Prevent MessageDispatch from recursing endlessly.
    assert ctx.messageMode != MessageMode::MAYBE_ENCRYPTED;
    self.set_input(self.start);
  }

  # One recursion to parse with the new ctx.messageMode setting.
  : MessageDispatch(ctx);
};

#-----------------------------------------------------------------------------
type EncryptedMessage = unit {
  len: uint32;
  : skip bytes &size=self.len;
};

#-----------------------------------------------------------------------------
type TlsForward = unit {
  # Just consume everything. This is hooked in ldap_zeek.spicy
  chunk: bytes &chunked &eod;
};

type KrbWrapToken = unit {
  # https://datatracker.ietf.org/doc/html/rfc4121#section-4.2.6.2

  # Number of bytes to expect *after* the payload.
  var trailer_ec: uint64;
  var header_ec: uint64;

  ctx_flags: bitfield(8) {
    send_by_acceptor: 0;
    sealed: 1;
    acceptor_subkey: 2;
  };
  filler: skip b"\xff";
  ec: uint16; # extra count
  rrc: uint16 { # right rotation count
    # Handle rrc == ec or rrc == 0.
    if ( self.rrc == self.ec ) {
      self.header_ec = self.ec;
    } else if ( self.rrc == 0 ) {
        self.trailer_ec = self.ec;
    } else {
      if ( ! self.ctx_flags.sealed )
        # If it's sealed, we'll consume until &eod anyhow
        # and ec/rrc shouldn't apply, otherwise, bail.
        throw "Unhandled rc %s and ec %s" % (self.ec, self.rrc);
    }
  }

  snd_seq: uint64;
  header_e: skip bytes &size=self.header_ec;
};

#-----------------------------------------------------------------------------
type SaslMsKrb5Stripper = unit(ctx: Ctx&) {
  # This is based on Wireshark output and example traffic we have. There's always
  # a 4 byte length field followed by the krb5_tok_id field in messages after
  # MS_KRB5 was selected. I haven't read enough specs to understand if it's
  # just this one case that works, or others could use the same stripping.
  var switch_size: uint64;

  len: uint32;
  krb5_tok_id: uint16;

  switch ( self.krb5_tok_id ) {
    0x0504 -> krb_wrap_token: KrbWrapToken;
    * -> : void;
  };

  : skip bytes &size=0 {
    self.switch_size = self.len - (self.offset() - 4);
    if ( self?.krb_wrap_token )
      self.switch_size -= self.krb_wrap_token.trailer_ec;
  }

  switch ( self?.krb_wrap_token && ! self.krb_wrap_token.ctx_flags.sealed ) {
    True -> : Message(ctx)[] &eod;
    * -> : skip bytes &eod;
  } &size=self.switch_size;

  # Consume the wrap token trailer, if any.
  trailer_e: skip bytes &size=self.krb_wrap_token.trailer_ec if (self?.krb_wrap_token);
};

#-----------------------------------------------------------------------------
public type Message = unit(ctx: Ctx&) {
  var messageID: int64;
  var opcode: ProtocolOpcode = ProtocolOpcode::Undef;
  var unsetResultDefault: Result;
  var result_: Result& = self.unsetResultDefault;
  var obj: string = "";
  var arg: string = "";
  var seqHeaderLen: uint64;
  var msgLen: uint64;
  var opLen: uint64;

  seqHeader: ASN1::ASN1Header &requires=($$.tag.class == ASN1::ASN1Class::Universal && $$.tag.type_ == ASN1::ASN1Type::Sequence) {
    self.msgLen = $$.len.len;
  }

  # Use offset() to determine how many bytes the seqHeader took. This
  # needs to be done after the seqHeader field hook.
  : void {
    self.seqHeaderLen = self.offset();
  }

  messageID_header: ASN1::ASN1Header &requires=($$.tag.class == ASN1::ASN1Class::Universal && $$.tag.type_ == ASN1::ASN1Type::Integer);
  : ASN1::ASN1Body(self.messageID_header, False) {
    self.messageID = $$.num_value;
  }

  protocolOp: ASN1::ASN1Header &requires=($$.tag.class == ASN1::ASN1Class::Application) {
    self.opcode = cast<ProtocolOpcode>(cast<uint8>($$.tag.type_));
    self.opLen = $$.len.len;
  }

  switch ( self.opcode ) {
    ProtocolOpcode::BIND_REQUEST            -> BIND_REQUEST:            BindRequest(self, ctx);
    ProtocolOpcode::BIND_RESPONSE           -> BIND_RESPONSE:           BindResponse(self, ctx);
    ProtocolOpcode::UNBIND_REQUEST          -> UNBIND_REQUEST:          UnbindRequest(self);
    ProtocolOpcode::SEARCH_REQUEST          -> SEARCH_REQUEST:          SearchRequest(self);
    ProtocolOpcode::SEARCH_RESULT_ENTRY     -> SEARCH_RESULT_ENTRY:     SearchResultEntry(self);
    ProtocolOpcode::SEARCH_RESULT_DONE      -> SEARCH_RESULT_DONE:      SearchResultDone(self);
    ProtocolOpcode::MODIFY_REQUEST          -> MODIFY_REQUEST:          ModifyRequest(self);
    ProtocolOpcode::MODIFY_RESPONSE         -> MODIFY_RESPONSE:         ModifyResponse(self);
    ProtocolOpcode::ADD_RESPONSE            -> ADD_RESPONSE:            AddResponse(self);
    ProtocolOpcode::DEL_REQUEST             -> DEL_REQUEST:             DelRequest(self);
    ProtocolOpcode::DEL_RESPONSE            -> DEL_RESPONSE:            DelResponse(self);
    ProtocolOpcode::MOD_DN_RESPONSE         -> MOD_DN_RESPONSE:         ModDNResponse(self);
    ProtocolOpcode::COMPARE_RESPONSE        -> COMPARE_RESPONSE:        CompareResponse(self);
    ProtocolOpcode::ABANDON_REQUEST         -> ABANDON_REQUEST:         AbandonRequest(self);

    # TODO: not yet implemented, redirect to NotImplemented because when we're
    # just commenting this out, it will stop processing LDAP Messages in this connection
    ProtocolOpcode::ADD_REQUEST             -> ADD_REQUEST:             NotImplemented(self);
    ProtocolOpcode::COMPARE_REQUEST         -> COMPARE_REQUEST:         NotImplemented(self);
    ProtocolOpcode::EXTENDED_REQUEST        -> EXTENDED_REQUEST:        ExtendedRequest(self, ctx);
    ProtocolOpcode::EXTENDED_RESPONSE       -> EXTENDED_RESPONSE:       ExtendedResponse(self, ctx);
    ProtocolOpcode::INTERMEDIATE_RESPONSE   -> INTERMEDIATE_RESPONSE:   NotImplemented(self);
    ProtocolOpcode::MOD_DN_REQUEST          -> MOD_DN_REQUEST:          NotImplemented(self);
    ProtocolOpcode::SEARCH_RESULT_REFERENCE -> SEARCH_RESULT_REFERENCE: NotImplemented(self);
  } &size=self.opLen;

  # Ensure some invariants hold after parsing the command.
  : void &requires=(self.offset() >= self.seqHeaderLen);
  : void &requires=(self.msgLen >= (self.offset() - self.seqHeaderLen));

  # Eat the controls field if it exists.
  : skip bytes &size=self.msgLen - (self.offset() - self.seqHeaderLen);
};

#-----------------------------------------------------------------------------
# Bind Operation
# https://tools.ietf.org/html/rfc4511#section-4.2

public type BindAuthType = enum {
  BIND_AUTH_SIMPLE = 0,
  BIND_AUTH_SASL = 3,

  # https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-adts/8b9dbfb2-5b6a-497a-a533-7e709cb9a982
  # 5.1.1.1.3 Sicily Authentication
  SICILY_PACKAGE_DISCOVERY = 9,
  SICILY_NEGOTIATE = 10,
  SICILY_RESPONSE = 11,
};

type GSS_SPNEGO_negTokenInit = unit {
  oidHeader: ASN1::ASN1Header &requires=($$.tag.class == ASN1::ASN1Class::Universal && $$.tag.type_ == ASN1::ASN1Type::ObjectIdentifier);
  oid: ASN1::ASN1ObjectIdentifier(self.oidHeader.len.len) &requires=(self.oid.oidstring == "1.3.6.1.5.5.2");

  # TODO: Parse the rest of negTokenInit.
  : skip bytes &eod;
};

# Peak into GSS-SPNEGO payload and ensure it is indeed GSS-SPNEGO,
# or GSS-SPNEGO with a NTMLSSP payload that starts with NTLMSSP.
type GSS_SPNEGO_Init = unit {
  # This is the optional octet string in SaslCredentials.
  credentialsHeader: ASN1::ASN1Header &requires=($$.tag.type_ == ASN1::ASN1Type::OctetString);

  # Now we either have the initial message as specified in RFC2743 or
  # a continuation from RFC4178, or a "NTMLSSP" signature.
  #
  #   60 -> APPLICATION [0]  https://datatracker.ietf.org/doc/html/rfc2743#page-81
  #   a1 -> CHOICE [1]       https://www.rfc-editor.org/rfc/rfc4178#section-4.2
  #   "NTMLSSP"              https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-nlmp/907f519d-6217-45b1-b421-dca10fc8af0d
  #
  switch {
    -> spnegoInitByte: uint8(0x60);
    -> spnegoChoiceByte: uint8(0xa1);
    -> ntlmSignature: skip b"NTLMSSP"; # Unsupported, should forward to child analyzer!
  };

  spnegoLen: skip ASN1::LengthType if (self?.spnegoInitByte || self?.spnegoChoiceByte);

  # Peak into the SPNEGO_negTokenInit
  spnegoInitial: skip GSS_SPNEGO_negTokenInit if (self?.spnegoInitByte);
};

type SaslCredentials = unit() {
  mechanism: ASN1::ASN1Message(False) &convert=$$.body.str_value;

  # Peak into GSS-SPNEGO payload if we have any.
  switch ( self.mechanism ) {
    "GSS-SPNEGO" -> gss_spnego: GSS_SPNEGO_Init;
    * -> : skip bytes &eod;
  };
};

type SicilyMessage = unit() {
  # Just ensure the signature matches. We could do more,
  # but it'd be better to forward to an NTLM analyzer.
  signature: skip b"NTLMSSP";
  var signature_decoded: string = "NTLMSSP";
};

type GSS_SPNEGO_Subsequent = unit {
  switch {
    -> spnegoChoiceByte: uint8(0xa1);
    -> ntmlSignature: skip b"NTLMSSP"; # Unsupported, should forward to NTLM!
  };

  spnegoChoiceLen: skip ASN1::LengthType if (self?.spnegoChoiceByte);
  negTokenResp: GSS_SPNEGO_negTokenResp if (self?.spnegoChoiceByte);
};

type GSS_SPNEGO_negTokenResp = unit {
  var accepted: bool;
  var supportedMech: ASN1::ASN1Message;
  var responseToken: optional<bytes>;

  # Parse the contained Sequence.
  seq: ASN1::ASN1Message(True) {
    for ( msg in $$.body.seq.submessages ) {
      # https://www.rfc-editor.org/rfc/rfc4178#section-4.2.2
      if ( msg.application_id == 0 ) {
        self.accepted = msg.application_data == b"\x0a\x01\x00";
      } else if ( msg.application_id == 1 ) {
        self.supportedMech = msg;
      } else if ( msg.application_id == 2 ) {
        self.responseToken = msg.application_data;
      } else if ( msg.application_id == 3 ) {
        # ignore mechListMec
      } else {
        throw "unhandled NegTokenResp id %s" % msg.application_id;
      }
    }
  }

  switch ( self?.supportedMech ) {
    True -> supportedMechOid: ASN1::ASN1Message(False) &convert=$$.body.str_value;
    * -> : void;
  } &parse-from=self.supportedMech.application_data;
};

# TODO(fox-ds): A helper unit for requests for which no handling has been implemented.
# Eventually all uses of this unit should be replaced with actual parsers so this unit can be removed.
type NotImplemented = unit(inout message: Message) {
  : skip bytes &eod;
};

type BindRequest = unit(inout message: Message, ctx: Ctx&) {
  version: ASN1::ASN1Message(True) &convert=$$.body.num_value;
  name: ASN1::ASN1Message(True) &convert=$$.body.str_value {
    message.obj = self.name;
  }
  var authType: BindAuthType = BindAuthType::Undef;
  var authData: bytes = b"";
  var simpleCreds: string = "";

  : ASN1::ASN1Message(True) {
    if ($$?.application_id) {
      self.authType = cast<BindAuthType>(cast<uint8>($$.application_id));
      self.authData = $$.application_data;
    }
  }

  if ( |self.authData| > 0 ) {
    switch ( self.authType ) {
      BindAuthType::BIND_AUTH_SIMPLE ->
        : void {
          self.simpleCreds = self.authData.decode();
          message.arg = self.simpleCreds;
        }

      BindAuthType::BIND_AUTH_SASL ->
        saslCreds: SaslCredentials {
          message.arg = self.saslCreds.mechanism;
          ctx.saslMechanism = self.saslCreds.mechanism;
        }

      BindAuthType::SICILY_NEGOTIATE, BindAuthType::SICILY_RESPONSE ->
        sicilyMessage: SicilyMessage {
          message.arg = self.sicilyMessage.signature_decoded;
        }

      * -> : void;
    } &parse-from=self.authData;
  };
} &requires=(self?.authType && (self.authType != BindAuthType::Undef));

type ServerSaslCreds = unit {
  serverSaslCreds: ASN1::ASN1Header &requires=($$.tag.class == ASN1::ASN1Class::ContextSpecific && $$.tag.type_ == ASN1::ASN1Type(7));
  payload: bytes &size=self.serverSaslCreds.len.len;
};

type BindResponse = unit(inout message: Message, ctx: Ctx&) {
  : Result {
    message.result_ = $$;

    # The SASL authentication was successful. We do not actually
    # know if the following messages are encrypted or not. This may be
    # mechanism and parameter specific. For example SCRAM-SHA512 or NTLM
    # will continue to be cleartext, while SRP or GSS-API would be encrypted.
    #
    # Switch messageMode into trial mode which is explored via MessageDispatch
    # and the MaybeEncrypted unit.
    #
    # Note, messageMode may be changed to something more specific like
    # MS_KRB5 below.
    if ( |ctx.saslMechanism| > 0 && $$.code == ResultCode::SUCCESS ) {
      ctx.messageMode = MessageMode::MAYBE_ENCRYPTED;
    }
  }

  # Try to parse serverSaslCreds if there's any input remaining. This
  # unit is parsed with &size, so &eod here works.
  #
  # Technically we should be able to tell from the ASN.1 structure
  # if the serverSaslCreds field exists or not. But, not sure we can
  # check if there's any bytes left at this point outside of passing
  # in the length and playing with offset().
  serverSaslCreds: ServerSaslCreds[] &eod;

  # If the client requested GSS-SPNEGO, try to parse the server's response
  # to switch message mode.
  gss_spnego: GSS_SPNEGO_Subsequent &parse-from=self.serverSaslCreds[0].payload
      if (ctx.saslMechanism == "GSS-SPNEGO" && |self.serverSaslCreds| > 0) {

    if ( $$?.negTokenResp ) {
      local token = $$.negTokenResp;
      if ( token.accepted && token?.supportedMechOid ) {
        if ( token.supportedMechOid == GSSAPI_MECH_MS_KRB5 && token.responseToken ) {
          ctx.messageMode = MessageMode::MS_KRB5;
        }
      }
    }
  }
};

#-----------------------------------------------------------------------------
# Unbind Operation
# https://tools.ietf.org/html/rfc4511#section-4.3

type UnbindRequest = unit(inout message: Message) {
  # this page intentionally left blank
};

#-----------------------------------------------------------------------------
# Search Operation
# https://tools.ietf.org/html/rfc4511#section-4.5

public type SearchScope = enum {
  SEARCH_BASE    = 0,
  SEARCH_SINGLE  = 1,
  SEARCH_TREE    = 2,
};

public type SearchDerefAlias = enum {
  DEREF_NEVER        = 0,
  DEREF_IN_SEARCHING = 1,
  DEREF_FINDING_BASE = 2,
  DEREF_ALWAYS       = 3,
};

type FilterType = enum {
  FILTER_AND     = 0,
  FILTER_OR      = 1,
  FILTER_NOT     = 2,
  FILTER_EQ      = 3,
  FILTER_SUBSTR  = 4,
  FILTER_GE      = 5,
  FILTER_LE      = 6,
  FILTER_PRESENT = 7,
  FILTER_APPROX  = 8,
  FILTER_EXT     = 9,
  FILTER_INVALID = 254,
};

public type AttributeSelection = unit {
  var attributes: vector<string>;

  # TODO: parse AttributeSelection as per
  # https://tools.ietf.org/html/rfc4511#section-4.5.1
  # and decide how deep that should be fleshed out.
  : ASN1::ASN1Message(True) {
    if (($$.head.tag.type_ == ASN1::ASN1Type::Sequence) &&
        ($$.body?.seq)) {
      for (i in $$.body.seq.submessages) {
        if (i.body?.str_value) {
          self.attributes.push_back(i.body.str_value);
        }
      }
    }
  }
};

type AttributeValueAssertion = unit {
  var desc: string = "";
  var val: string = "";

  : ASN1::ASN1Message(True) {
    if (($$.head.tag.type_ == ASN1::ASN1Type::Sequence) &&
        ($$.body?.seq) &&
        (|$$.body.seq.submessages| >= 2)) {
      if ($$.body.seq.submessages[0].body?.str_value) {
        self.desc = $$.body.seq.submessages[0].body.str_value;
      }
      if ($$.body.seq.submessages[1].body?.str_value) {
        self.val = $$.body.seq.submessages[1].body.str_value;
      }
    }
  }
};

# An AND or OR search filter can consist of many sub-searchfilters, so we try to parse these
type ParseNestedAndOr = unit {
    searchfilters: SearchFilter[] &eod;
};

type ParseNestedNot = unit {
    searchfilter: SearchFilter;
};

# Helper functions to properly format some custom data structures

public function utf16_guid_to_hex_repr(bts: bytes) : string {
  # Rather ugly workaround to pretty-print the CLDAP DomainGuid UTF16-LE encoded string
  # in the same format as Wireshark (aabbccdd-eeff-gghh-iijj-kkllmmnnoopp)

  # We need to have exactly 16 bytes...
  if ( |bts| != 16 ) {
    # ... and otherwise just return an error code
    return "GUID_FORMAT_FAILED";
  }

  local ret = "";
  for ( i in [[3, 2, 1, 0], [5, 4], [7, 6], [8, 9],  [10, 11, 12, 13, 14, 15]] ) {
    for ( j in i ) {
      local bt: uint8 = *bts.at(j);
      ret = ret + "%02x" % bt;
      if ( j in [0, 4, 6, 9] ) {
        ret = ret + "-";
      }
    }
  }
  return ret;
}

public function bytes_sid_to_hex_repr(bts: bytes) : string {
  local ret = "";
  local cnt = 0;

  while ( cnt < |bts| ) {
    local bt: uint8 = *bts.at(cnt);
    ret = ret + "%02x" % bt;

    if ( cnt < |bts|-1 ) {
      ret = ret + ":";
    }
    cnt += 1;
  }
  return ret;
}

public function bytes_sid_to_SID_repr(bts: bytes) : string {
  # Example: SID -> S-1-5-21-1153942841-488947194-1912431946

  # Needs to be exactly 24 bytes
  if ( |bts| != 24 ) {
    # ... and otherwise just return an error code
    return "SID_FORMAT_FAILED";
  }

  local ret = "S-";
  local cnt = 0;

  # Mixed little and big endian, so turn everything to big endian first...
  # Byte 1 seems to be skipped when parsing the SID
  for ( i in [[0], [2, 3, 4, 5, 6, 7], [11, 10, 9, 8], [15, 14, 13, 12], [19, 18, 17, 16], [23, 22, 21, 20]] ) {
    local dec_val_rep: bytes = b"";
    for ( j in i ) {
      local bt: uint8 = *bts.at(j);
      dec_val_rep += bt;
      cnt += 1;
    }

    # ... so we can represent this integer value in big endian
    ret = ret + "%u" % dec_val_rep.to_uint(spicy::ByteOrder::Big);

    # Only print the dash when we're not at the end
    if ( cnt < 23 ) {
      ret = ret + "-";
    }
  }
  return ret;
}


public function uint32_to_hex_repr(bts: bytes) : string {
  # Needs to be exactly 4 bytes
  if ( |bts| != 4 ) {
    # ... and otherwise just return an error code
    return "HEX_FORMAT_FAILED";
  }

  # Workaround to print the hex value of an uint32, prepended with '0x'
  local ret = "0x";
  for ( i in [3, 2, 1, 0] ) {
    local bt: uint8 = *bts.at(i);
    ret = ret + "%02x" % bt;
  }
  return ret;
}

# Helper to compute a string representation of a `SearchFilter`.
public function string_representation(search_filter: SearchFilter): string {
    local repr: string;

    switch ( local fType = search_filter.filterType ) {
      # The NOT, AND and OR filter types are trees and may hold many leaf nodes. So recursively get
      # the stringPresentations for the leaf nodes and add them all in one final statement.

      case FilterType::FILTER_NOT: {
        repr = "(!%s)" % search_filter.FILTER_NOT.searchfilter.stringRepresentation;
      }

      case FilterType::FILTER_AND, FilterType::FILTER_OR: {
        local nestedObj: ParseNestedAndOr;
        local printChar = "";

        if ( fType == FilterType::FILTER_AND ) {
          printChar = "&";
          nestedObj = search_filter.FILTER_AND;
        } else {
          printChar = "|";
          nestedObj = search_filter.FILTER_OR;
        }

        # Build the nested AND/OR statement in this loop. When we encounter the first element,
        # we open the statement. At the second element, we close our first complete statement. For every
        # following statement, we extend the AND/OR statement by wrapping it around the already completed
        # statement. Although it is also valid to not do this wrapping, which is logically equivalent, e.g:
        #
        #           (1)                         (2)
        #   (?(a=b)(c=d)(e=f))    vs   (?(?(a=b)(c=d))(e=f))
        #
        # the latter version is also shown by Wireshark. So although the parsed structure actually represents
        # version (1) of the query, we now choose to print version (2). If this is not desirable, swap the code
        # for the following:
        #
        # # Construct the nested structure, like (1)
        # for ( SF in nestedObj.searchfilters ) {
        #    self.stringRepresentation = self.stringRepresentation + SF.stringRepresentation
        # }
        # # Close it with brackets and put the correct printChar for AND/OR in the statement
        # self.stringRepresentation =  "(%s" % printChar + self.stringRepresentation + ")";
        #

        local i = 0;
        for ( searchFilter in nestedObj.searchfilters ) {
          switch ( i ) {
            case 0: {
              repr = "(%s%s%s" % (
                  printChar,
                  searchFilter.stringRepresentation,
                  # If we have exactly one element immediately close the statement since we are done.
                  |nestedObj.searchfilters| == 1 ? ")" : ""
              );
            }
            case 1: {
              repr = repr + searchFilter.stringRepresentation + ")";
            }
            default: {
              repr = "(%s" % printChar + repr + searchFilter.stringRepresentation + ")";
            }
          }
          i += 1;
        }
      }

    # The following FilterTypes are leaf nodes and can thus be represented in a statement

      case FilterType::FILTER_EXT: {
        # For extended search filters the meaning of the individual fields in
        # `DecodedAttributeValue` is slightly different.
        repr = "(%s:%s:=%s)" % (search_filter.FILTER_EXT.assertionValueDecoded,
                                search_filter.FILTER_EXT.attributeDesc.decode(),
                                search_filter.FILTER_EXT.matchValue);
      }
      case FilterType::FILTER_APPROX: {
        repr = "(%s~=%s)" % (search_filter.FILTER_APPROX.attributeDesc.decode(),
                             search_filter.FILTER_APPROX.assertionValueDecoded);
      }
      case FilterType::FILTER_EQ: {
        repr = "(%s=%s)" % (search_filter.FILTER_EQ.attributeDesc.decode(),
                            search_filter.FILTER_EQ.assertionValueDecoded);
      }
      case FilterType::FILTER_GE: {
        repr = "(%s>=%s)" % (search_filter.FILTER_GE.attributeDesc.decode(),
                             search_filter.FILTER_GE.assertionValueDecoded);
      }
      case FilterType::FILTER_LE: {
        repr = "(%s<=%s)" % (search_filter.FILTER_LE.attributeDesc.decode(),
                             search_filter.FILTER_LE.assertionValueDecoded);
      }
      case FilterType::FILTER_SUBSTR: {
        local anys: string = "";
        if ( |search_filter.FILTER_SUBSTR.anys| > 0 )
          anys = b"*".join(search_filter.FILTER_SUBSTR.anys).decode() + "*";

        repr  = "(%s=%s*%s%s)" % (search_filter.FILTER_SUBSTR.attributeDesc.decode(),
                                  search_filter.FILTER_SUBSTR.initial,
                                  anys,
                                  search_filter.FILTER_SUBSTR.final);
      }
      case FilterType::FILTER_PRESENT: {
        repr = "(%s=*)" % search_filter.FILTER_PRESENT;
      }
    }

    return repr;
}

# Represents an (extended) key-value pair present in SearchFilters
type DecodedAttributeValue = unit(fType: FilterType) {
  var assertionValueDecoded: string = "";

  : uint8;
  attributeDesc_len: uint8;
  attributeDesc: bytes &size=self.attributeDesc_len;

  : uint8;
  assertionValue_len: uint8;
  assertionValue: bytes &size=self.assertionValue_len;

  # Only for the FILTER_EXT type, parse extra fields
  : uint8 if ( fType == FilterType::FILTER_EXT );
  matchValue_len: uint8 if( fType == FilterType::FILTER_EXT );
  matchValue: bytes &size=self.matchValue_len if ( fType == FilterType::FILTER_EXT );

  on %done {
    switch ( self.attributeDesc ) {
      # Special parsing required for some CLDAP attributes,
      # see https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-adts/895a7744-aff3-4f64-bcfa-f8c05915d2e9

      case b"DomainGuid": {
        self.assertionValueDecoded = utf16_guid_to_hex_repr(self.assertionValue);
      }

      case b"objectSid", b"AAC": {
        self.assertionValueDecoded = bytes_sid_to_hex_repr(self.assertionValue);
      }

      case b"DomainSid": {
        self.assertionValueDecoded = bytes_sid_to_SID_repr(self.assertionValue);
      }

      case b"NtVer": {
        self.assertionValueDecoded = uint32_to_hex_repr(self.assertionValue);
      }

      # By default, decode with UTF-8
      default: {
        self.assertionValueDecoded = self.assertionValue.decode();
      }
    }
  }
};

type SubstringFilter = unit {
  var initial: string;
  var final: string;
  var anys: vector<string>;

  : uint8;  # filter tag
  attributeDesc_len: uint8;
  attributeDesc: bytes &size=self.attributeDesc_len;

  # Crunch through the sequence/choice of substrings.
  #
  # https://datatracker.ietf.org/doc/html/rfc4511#section-4.5.1
  header: ASN1::ASN1Header;
  : ASN1::ASN1Message(False)[] &size=self.header.len.len foreach {
    local data = $$.application_data.decode();
    if ( $$.application_id == 0 ) {
      self.initial = data;
    } else if ( $$.application_id == 1 ) {
      self.anys.push_back(data);
    } else if ( $$.application_id == 2 ) {
      self.final = data;
    } else {
      throw "invalid substring choice %s" % $$.application_id;
    }
  }
};

type SearchFilter = unit {
  var filterType: FilterType = FilterType::Undef;
  var filterBytes: bytes = b"";
  var filterLen: uint64 = 0;
  var stringRepresentation: string = "";

  : ASN1::ASN1Message(True) {
    if ($$?.application_id) {
      self.filterType = cast<FilterType>(cast<uint8>($$.application_id));
      self.filterBytes = $$.application_data;
      self.filterLen = $$.head.len.len;
    } else {
      self.filterType = FilterType::FILTER_INVALID;
    }
  }

  switch ( self.filterType ) {

    # FilterTypes that hold one or more SearchFilters inside them

    FilterType::FILTER_AND     -> FILTER_AND:     ParseNestedAndOr()
                                                  &parse-from=self.filterBytes;
    FilterType::FILTER_OR      -> FILTER_OR:      ParseNestedAndOr()
                                                  &parse-from=self.filterBytes;
    FilterType::FILTER_NOT     -> FILTER_NOT:     ParseNestedNot()
                                                  &parse-from=self.filterBytes;

    # FilterTypes that we can actually convert to a string

    FilterType::FILTER_EQ      -> FILTER_EQ:      DecodedAttributeValue(FilterType::FILTER_EQ)
                                                  &parse-from=self.filterBytes;
    FilterType::FILTER_SUBSTR  -> FILTER_SUBSTR:  SubstringFilter
                                                  &parse-from=self.filterBytes;
    FilterType::FILTER_GE      -> FILTER_GE:      DecodedAttributeValue(FilterType::FILTER_GE)
                                                  &parse-from=self.filterBytes;
    FilterType::FILTER_LE      -> FILTER_LE:      DecodedAttributeValue(FilterType::FILTER_LE)
                                                  &parse-from=self.filterBytes;
    FilterType::FILTER_APPROX  -> FILTER_APPROX:  DecodedAttributeValue(FilterType::FILTER_APPROX)
                                                  &parse-from=self.filterBytes;
    FilterType::FILTER_EXT     -> FILTER_EXT:     DecodedAttributeValue(FilterType::FILTER_EXT)
                                                  &parse-from=self.filterBytes;
    FilterType::FILTER_PRESENT -> FILTER_PRESENT: ASN1::ASN1OctetString(self.filterLen, False)
                                                  &convert=$$.value.decode(spicy::Charset::ASCII)
                                                  &parse-from=self.filterBytes;
  };

  # So when you're done with recursively parsing the filters, we can now leverage the tree structure to
  # recursively get the stringRepresentations for those leaves, which are SearchFilters

  on %done {
    self.stringRepresentation = string_representation(self);
  }

  on %error {
    self.stringRepresentation = "FILTER_PARSING_ERROR";
  }

};

public type SearchRequest = unit(inout message: Message)  {
  baseObject: ASN1::ASN1Message(True) &convert=$$.body.str_value {
    message.obj = self.baseObject;
  }
  scope: ASN1::ASN1Message(True) &convert=cast<SearchScope>(cast<uint8>($$.body.num_value))
                                 &default=SearchScope::Undef {
    message.arg = "%s" % self.scope;
  }
  deref: ASN1::ASN1Message(True) &convert=cast<SearchDerefAlias>(cast<uint8>($$.body.num_value))
                                 &default=SearchDerefAlias::Undef;
  sizeLimit: ASN1::ASN1Message(True) &convert=$$.body.num_value &default=0;
  timeLimit: ASN1::ASN1Message(True) &convert=$$.body.num_value &default=0;
  typesOnly: ASN1::ASN1Message(True) &convert=$$.body.bool_value &default=False;
  filter: SearchFilter &convert=$$.stringRepresentation;
  attributes: AttributeSelection &convert=$$.attributes;
};

type SearchResultEntry = unit(inout message: Message) {
  objectName: ASN1::ASN1Message(True) &convert=$$.body.str_value {
    message.obj = self.objectName;
  }
  # TODO: if we want to descend down into PartialAttributeList
  attributes: ASN1::ASN1Message(True);
};

type SearchResultDone = unit(inout message: Message) {
  : Result {
    message.result_ = $$;
  }
};

# TODO: implement SearchResultReference
# type SearchResultReference = unit(inout message: Message) {
#
# };

#-----------------------------------------------------------------------------
# Modify Operation
# https://tools.ietf.org/html/rfc4511#section-4.6

type ModifyRequest = unit(inout message: Message) {
  objectName: ASN1::ASN1Message(True) &convert=$$.body.str_value {
    message.obj = self.objectName;
  }

  # TODO: parse changes
};

type ModifyResponse = unit(inout message: Message) {
  : Result {
    message.result_ = $$;
  }
};

#-----------------------------------------------------------------------------
# Add Operation
# https://tools.ietf.org/html/rfc4511#section-4.7

# TODO: implement AddRequest
# type AddRequest = unit(inout message: Message) {
#
#
# };

type AddResponse = unit(inout message: Message) {
  : Result {
    message.result_ = $$;
  }
};

#-----------------------------------------------------------------------------
# Delete Operation
# https://tools.ietf.org/html/rfc4511#section-4.8

type DelRequest = unit(inout message: Message) {
  objectName: ASN1::ASN1Message(True) &convert=$$.body.str_value {
    message.obj = self.objectName;
  }
};

type DelResponse = unit(inout message: Message) {
  : Result {
    message.result_ = $$;
  }
};

#-----------------------------------------------------------------------------
# Modify DN Operation
# https://tools.ietf.org/html/rfc4511#section-4.8

# TODO: implement ModDNRequest
# type ModDNRequest = unit(inout message: Message) {
#
# };

type ModDNResponse = unit(inout message: Message) {
  : Result {
    message.result_ = $$;
  }
};

#-----------------------------------------------------------------------------
# Compare Operation
# https://tools.ietf.org/html/rfc4511#section-4.10

# TODO: implement CompareRequest
# type CompareRequest = unit(inout message: Message) {
#
# };

type CompareResponse = unit(inout message: Message) {
  : Result {
    message.result_ = $$;
  }
};

#-----------------------------------------------------------------------------
# Abandon Operation
# https://tools.ietf.org/html/rfc4511#section-4.11

type AbandonRequest = unit(inout message: Message) {
  messageID: ASN1::ASN1Message(True) &convert=$$.body.num_value {
    message.obj = "%d" % (self.messageID);
  }
};

#-----------------------------------------------------------------------------
# Extended Operation
# https://tools.ietf.org/html/rfc4511#section-4.12
type ExtendedRequest = unit(inout message: Message, ctx: Ctx&) {
  var requestValue: bytes;
  header: ASN1::ASN1Header &requires=($$.tag.class == ASN1::ASN1Class::ContextSpecific);
  requestName: bytes &size=self.header.len.len &convert=$$.decode(spicy::Charset::ASCII) {
    message.obj = $$;
  }

  # If there's more byte to parse, it's the requestValue.
  : ASN1::ASN1Message(False)
      &requires=($$.head.tag.class == ASN1::ASN1Class::ContextSpecific)
      if ( message.opLen > self.offset() ) {

    self.requestValue = $$.application_data;
  }

  on %done {
    # Did the client request StartTLS?
    #
    # https://datatracker.ietf.org/doc/html/rfc4511#section-4.14.1
    if ( self.requestName == "1.3.6.1.4.1.1466.20037" )
      ctx.startTlsRequested = True;
  }
};

#-----------------------------------------------------------------------------
type ExtendedResponseEntry = unit(inout r: ExtendedResponse) {
  : ASN1::ASN1Message(False) &requires=($$.head.tag.class == ASN1::ASN1Class::ContextSpecific) {
    if ( $$.head.tag.type_ == ASN1::ASN1Type(10) )
      r.responseName = $$.application_data;
    else if ( $$.head.tag.type_ == ASN1::ASN1Type(11) )
      r.responseValue = $$.application_data;
    else
      throw "Unhandled extended response tag %s" % $$.head.tag;
  }
};

#-----------------------------------------------------------------------------
type ExtendedResponse = unit(inout message: Message, ctx: Ctx&) {
  var responseName: bytes;
  var responseValue: bytes;
  : Result {
    message.result_ = $$;
  }

  # Try to parse two ASN1 entries if there are bytes left in the unit.
  # Both are optional and identified by context specific tagging.
  : ExtendedResponseEntry(self) if ( message.opLen > self.offset() );
  : ExtendedResponseEntry(self) if ( message.opLen > self.offset() );

  on %done {
    # Client had requested StartTLS and it was successful? Switch to SSL.
    if ( ctx.startTlsRequested && message.result_.code == ResultCode::SUCCESS )
      ctx.messageMode = MessageMode::TLS;
  }
};

#-----------------------------------------------------------------------------
# IntermediateResponse Message
# https://tools.ietf.org/html/rfc4511#section-4.13

# TODO: implement IntermediateResponse
# type IntermediateResponse = unit(inout message: Message) {
#
# };

on LDAP::Message::%done {
  spicy::accept_input();
}
