# -*- coding: binary -*-

require 'rasn1'

module Msf
  class Exploit
    class Remote
      module Kerberos
        module Client
          # Methods for interacting with Kerberos's PKINIT extension for obtaining a
          # TGT from a certificate
          #
          # https://www.rfc-editor.org/rfc/rfc4556
          # https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-pkca/d0cf1763-3541-4008-a75f-a577fa5e8c5b
          module Pkinit
            # Builds a Diffie Helman object with parameters set up
            #
            # @return [OpenSSL::PKey::DH, string] The Diffie Hellman object, and a random client nonce
            def build_dh
              # When using the Diffie-Hellman key agreement method, implementations MUST support Oakley 1024-bit Modular
              # Exponential (MODP) well-known group 2 RFC2412
              # Kerberos spec: https://www.rfc-editor.org/rfc/rfc4556
              # Value: https://www.rfc-editor.org/rfc/rfc2412#appendix-E.2
              prime_modulus = 0 # built 256 bits at a time
              prime_modulus |= 0xffffffffffffffffc90fdaa22168c234c4c6628b80dc1cd129024e088a67cc74 << (256 * 3)
              prime_modulus |= 0x020bbea63b139b22514a08798e3404ddef9519b3cd3a431b302b0a6df25f1437 << (256 * 2)
              prime_modulus |= 0x4fe1356d6d51c245e485b576625e7ec6f44c42e9a637ed6b0bff5cb6f406b7ed << (256 * 1)
              prime_modulus |= 0xee386bfb5a899fa5ae9f24117c4b1fe649286651ece65381ffffffffffffffff
              dh = OpenSSL::PKey::DH.new(
                OpenSSL::ASN1::Sequence([
                  OpenSSL::ASN1::Integer(prime_modulus),
                  OpenSSL::ASN1::Integer(2)
                ]).to_der
              )
              if OpenSSL::PKey.respond_to?(:generate_key)
                # OpenSSL v3.x path
                # see:
                #  * https://github.com/rapid7/metasploit-framework/pull/17308
                #  * https://ruby-doc.org/stdlib-3.1.0/libdoc/openssl/rdoc/OpenSSL/PKey/DH.html#method-i-generate_key-21
                dh = OpenSSL::PKey.generate_key(dh)
              else
                dh.generate_key!
              end

              dh_nonce = SecureRandom.random_bytes(32)
              [dh, dh_nonce]
            end

            # Extracts the user and realm from a certificate, deferring to
            # the provided values if they are not nil.
            #
            # @param certificate [OpenSSL::X509::Certificate]
            # @param username [String] A default value for username. A warning is presented if this is not in the certificate.
            # @param realm [String] A default value for realm. A warning is presented if this is not in the certificate.
            # @return [Array<String>] A tuple of the username and realm retrieved from the certificate, or parameters provided
            # @raise [ArgumentError] If the certificate contains a corrupted SAN
            # @raise [ArgumentError] If a username is provided without also providing a realm; or vice versa
            def extract_user_and_realm(certificate, username, realm)
              raise ArgumentError, 'Must provide username if providing realm' if username.nil? && !realm.nil?
              raise ArgumentError, 'Must provide realm if providing username' if realm.nil? && !username.nil?

              results = []
              asn_san_seq = []

              # MS's SAN extension isn't handled nicely by OpenSSL, so we need to read it ourselves
              # https://manas.tech/blog/2013/01/29/extracting-subject-alternative-name-from-microsoft-authentication-client-certificates/
              certificate.extensions.select { |ext| ext.oid == 'subjectAltName' }.each do |san_extension|
                begin
                  asn_san = OpenSSL::ASN1.decode(san_extension)
                  asn_san_value = asn_san.value.find { |value| value.is_a? OpenSSL::ASN1::OctetString }

                  if asn_san_value.nil?
                    raise ArgumentError, 'Invalid certificate provided: unable to decode SAN'
                  end

                  asn_san_seq = OpenSSL::ASN1.decode(asn_san_value.value)
                rescue OpenSSL::ASN1::ASN1Error
                  raise ArgumentError, 'Invalid certificate provided: unable to decode SAN'
                end

                asn_san_seq.each do |san_entry|
                  if san_entry.tag == 0 # x509.OtherName
                    key = san_entry.value[0]&.value
                    next if key != 'msUPN' # Principal Name

                    principal = san_entry.value[1].value[0].value
                    parts = principal.split('@')
                    if parts.length == 1
                      user = principal
                      domain = ''
                    else
                      user = parts[0..-2].join('@')
                      domain = parts[-1]
                    end
                  elsif san_entry.tag == 2 # dNSName
                    parts = san_entry.value.split('.')
                    if parts.length == 1
                      user = san_entry.value  # Corrected to extract string value
                      domain = ''
                    else
                      user = parts[0] + '$'
                      domain = parts[1..].join('.')
                    end
                  else
                    next
                  end

                  results.append([user, domain])
                end
              end

              unless realm.nil? # and also username, since it's both or neither
                normalized_results = results.map do |pair|
                  pair.map do |value|
                    if value.is_a?(String)
                      value.downcase
                    elsif value.is_a?(OpenSSL::ASN1::ASN1Data) && value.respond_to?(:value)
                      val = value.value
                      val.is_a?(String) ? val.downcase : val.to_s.downcase
                    else
                      value.to_s.downcase
                    end
                  end
                end

                unless normalized_results.include?([username.downcase, realm.downcase])
                  print_warning("Warning: Provided principal and realm (#{username}@#{realm}) do not match entries in certificate:")
                  results.each do |cert_username, cert_realm|
                    print_warning("  * #{cert_username}@#{cert_realm}")
                  end
                end

                return [username, realm]
              end

              # No override was provided, so hopefully we only extracted one value from the certificate
              if results.length == 1
                return results[0]
              else
                raise ArgumentError, "Failed to retrieve Principal from certificate (contained #{results.length} SAN entries). Provide an override user and domain."
              end
            end

            # Transform a key into a key of a certain size, using the k-truncate algorithm
            # described in https://www.rfc-editor.org/rfc/rfc4556#section-3.2.3.1
            #
            # @param data [String] The full key to transform
            # @param etype [Integer] The encryption type, from Rex::Proto::Kerberos::Crypto::Encryption
            # @return [String] The truncated key
            def k_truncate(data, etype)
              if etype == Rex::Proto::Kerberos::Crypto::Encryption::AES256
                keysize = 32
              elsif etype == Rex::Proto::Kerberos::Crypto::Encryption::AES128
                keysize = 16
              else
                # This is unsupported per the spec
                raise Rex::Proto::Kerberos::Model::Error::KerberosEncryptionNotSupported.new("Unsupported DH Key exchange encryption type #{etype}", encryption_type: etype)
              end

              result = ''
              x = 0
              while result.length < keysize
                digest = Digest::SHA1.digest(x.chr + data)
                if result.length + digest.length > keysize
                  result += digest[0..(keysize - result.length - 1)] # Just take the first few bytes until we reach the desired length
                  return result
                end
                result += digest
                x += 1
              end

              result
            end

            # Given all the Diffie Hellman parameters and response from the server,
            # calculate the shared key using the steps described in
            # https://www.rfc-editor.org/rfc/rfc4556#section-3.2.3.1
            #
            # @param pa_pk_as_rep [Rex::Proto::Kerberos::Model::PreAuthPkAsRep] The PA_DATA response from the server containing the server's public key
            # @param dh [OpenSSL::PKey::DH, string] The Diffie Hellman object
            # @param dh_nonce [String] The random client nonce we sent to the server
            # @param etype [Integer] The encryption type, from Rex::Proto::Kerberos::Crypto::Encryption
            # @return [String] The calculated shared key
            def calculate_shared_key(pa_pk_as_rep, dh, dh_nonce, etype)
              dh_rep_info = pa_pk_as_rep.dh_rep_info
              signed_data = dh_rep_info.signed_data
              dh_key_info = signed_data[:encap_content_info].econtent
              server_public_key = RASN1::Types::Integer.parse(dh_key_info[:subject_public_key].value).value
              shared_key = dh.compute_key(server_public_key.to_bn)
              server_nonce = pa_pk_as_rep[:server_dh_nonce].value
              full_key = shared_key + dh_nonce + server_nonce
              k_truncate(full_key, etype)
            end

            # Build a PreAuth data entry structure for negotiating a shared DH key with the server
            #
            # @param pfx [OpenSSL::PKCS12] A PKCS12-encoded certificate
            # @param dh [OpenSSL::PKey::DH, string] The Diffie Hellman object
            # @param dh_nonce [String] The random client nonce we sent to the server
            # @param request_body [Rex::Proto::Kerberos::Model::KdcRequest] The request body accompanying this PreAuth entry
            # @param opts [Hash] Options to override default values for certain PreAuth entry fields
            # @return  [Rex::Proto::Kerberos::Model::PreAuthDataEntry] The constructed PreAuth data entry request
            def build_pa_pk_as_req(pfx, dh, dh_nonce, request_body, opts)
              certificate = pfx.certificate
              now_time = kerberos_time
              now_ctime = now_time.round
              ctime = opts.fetch(:ctime) { now_ctime }
              cusec = opts.fetch(:cusec) { now_time&.usec || 0 }
              nonce = opts.fetch(:nonce) { rand(1 << 31) }
              data = request_body.encode
              checksum = Digest::SHA1.digest(data)
              pub_key_encoded = RASN1::Types::Integer.new(value: dh.pub_key.to_i).to_der
              auth_pack = Rex::Proto::Kerberos::Model::Pkinit::AuthPack.new(
                pk_authenticator: {
                  cusec: cusec,
                  ctime: ctime,
                  nonce: nonce,
                  pa_checksum: checksum
                },
                client_public_value: {
                  algorithm: {
                    algorithm: Rex::Proto::Kerberos::Model::OID::DiffieHellman, # Diffie-Hellman
                    parameters: Rex::Proto::Kerberos::Model::Pkinit::DomainParameters.new(
                      p: dh.p.to_i,
                      g: dh.g.to_i,
                      q: 0
                    )
                  },
                  subject_public_key: pub_key_encoded
                },
                client_dh_nonce: RASN1::Types::OctetString.new(value: dh_nonce)
              )


              auth_pack[:client_public_value][:subject_public_key].bit_length = pub_key_encoded.length * 8


              signed_auth_pack = sign_auth_pack(auth_pack, pfx.key, certificate)


              pa_as_req = Rex::Proto::Kerberos::Model::PreAuthPkAsReq.new


              pa_as_req.signed_auth_pack = signed_auth_pack


              Rex::Proto::Kerberos::Model::PreAuthDataEntry.new(type: Rex::Proto::Kerberos::Model::PreAuthType::PA_PK_AS_REQ,
                                                                  value: pa_as_req.to_der)
            end

            # Calculate the cryptographic signatures over the AuthPack, and create the appropriate
            # ASN.1-encoded structure, per https://www.rfc-editor.org/rfc/rfc4556#section-3.2.1
            #
            # @param auth_pack [Rex::Proto::Kerberos::Model::Pkinit::AuthPack] The AuthPack to sign
            # @param key [OpenSSL::PKey] The private key to digitally sign the data
            # @param dh [OpenSSL::X509::Certificate] The certificate associated with the private key
            # @return [Rex::Proto::CryptoAsn1::Cms::ContentInfo] The signed AuthPack
            def sign_auth_pack(auth_pack, key, certificate)
              signer_info = Rex::Proto::CryptoAsn1::Cms::SignerInfo.new(
                version: 1,
                sid: {
                  issuer: certificate.issuer,
                  serial_number: certificate.serial.to_i
                },
                digest_algorithm: {
                  algorithm: Rex::Proto::Kerberos::Model::OID::SHA1
                },
                signed_attrs: [
                  {
                    attribute_type: Rex::Proto::Kerberos::Model::OID::ContentType,
                    attribute_values: [RASN1::Types::Any.new(value: RASN1::Types::ObjectId.new(value: Rex::Proto::Kerberos::Model::OID::PkinitAuthData))]
                  },
                  {
                    attribute_type: Rex::Proto::Kerberos::Model::OID::MessageDigest,
                    attribute_values: [RASN1::Types::Any.new(value: RASN1::Types::OctetString.new(value: Digest::SHA1.digest(auth_pack.to_der)))]
                  }
                ],
                signature_algorithm: {
                  algorithm: Rex::Proto::Kerberos::Model::OID::RSAWithSHA1
                }
              )
              data = RASN1::Types::Set.new(value: signer_info[:signed_attrs].value).to_der
              signature = key.sign(OpenSSL::Digest.new('SHA1'), data)

              signer_info[:signature] = signature

              signed_data = Rex::Proto::CryptoAsn1::Cms::SignedData.new(
                version: 3,
                digest_algorithms: [
                  {
                    algorithm: Rex::Proto::Kerberos::Model::OID::SHA1
                  }
                ],
                encap_content_info: {
                  econtent_type: Rex::Proto::Kerberos::Model::OID::PkinitAuthData,
                  econtent: auth_pack.to_der
                },
                certificates: [{ openssl_certificate: certificate }],
                signer_infos: [signer_info]
              )

              Rex::Proto::CryptoAsn1::Cms::ContentInfo.new(
                content_type: Rex::Proto::Kerberos::Model::OID::SignedData,
                data: signed_data
              )
            end
          end
        end
      end
    end
  end
end
