###
#
# This mixin provides methods to add, delete and lookup computer accounts via MS-SAMR
#
# -*- coding: binary -*-

module Msf

module Exploit::Remote::MsSamr

  include Msf::Exploit::Remote::SMB::Client::Authenticated
  include Msf::Auxiliary::Report

  class MsSamrError < StandardError; end
  class MsSamrConnectionError < MsSamrError; end
  class MsSamrAuthenticationError < MsSamrError; end
  class MsSamrNotFoundError < MsSamrError; end
  class MsSamrUnexpectedReplyError < MsSamrError; end
  class MsSamrUnknownError < MsSamrError; end
  class MsSamrBadConfigError < MsSamrError; end

  ComputerInfo = Struct.new(:name, :password)
  SamrConnection = Struct.new(:samr, :server_handle, :domain_handle, :domain_name)

  def initialize(info = {})
    super

    register_options([
      OptString.new('COMPUTER_NAME', [ false, 'The computer name' ]),
      OptString.new('COMPUTER_PASSWORD', [ false, 'The password for the new computer' ]),
    ], Msf::Exploit::Remote::MsSamr)
  end

  def add_computer(opts = {})
    tree = opts[:tree] || connect_ipc

    samr_con = connect_samr(tree)

    computer_name = opts[:computer_name] || datastore['COMPUTER_NAME']
    if computer_name.blank?
      computer_name = random_hostname
      4.downto(0) do |attempt|
        break if samr_con.samr.samr_lookup_names_in_domain(
          domain_handle: samr_con.domain_handle,
          names: [ computer_name ]
        ).nil?

        computer_name = random_hostname
        raise MsSamrBadConfigError, 'Could not find an unused computer name.' if attempt == 0
      end
    else
      if samr_con.samr.samr_lookup_names_in_domain(domain_handle: samr_con.domain_handle, names: [ computer_name ])
        raise MsSamrBadConfigError, 'The specified computer name already exists.'
      end
    end

    result = samr_con.samr.samr_create_user2_in_domain(
      domain_handle: samr_con.domain_handle,
      name: computer_name,
      account_type: RubySMB::Dcerpc::Samr::USER_WORKSTATION_TRUST_ACCOUNT,
      desired_access: RubySMB::Dcerpc::Samr::USER_FORCE_PASSWORD_CHANGE | RubySMB::Dcerpc::Samr::MAXIMUM_ALLOWED
    )

    user_handle = result[:user_handle]
    if datastore['COMPUTER_PASSWORD'].blank?
      computer_password = Rex::Text.rand_text_alphanumeric(32)
    else
      computer_password = datastore['COMPUTER_PASSWORD']
    end

    user_info = RubySMB::Dcerpc::Samr::SamprUserInfoBuffer.new(
      tag: RubySMB::Dcerpc::Samr::USER_INTERNAL4_INFORMATION_NEW,
      member: RubySMB::Dcerpc::Samr::SamprUserInternal4InformationNew.new(
        i1: {
          password_expired: 1,
          which_fields: RubySMB::Dcerpc::Samr::USER_ALL_NTPASSWORDPRESENT | RubySMB::Dcerpc::Samr::USER_ALL_PASSWORDEXPIRED
        },
        user_password: {
          buffer: RubySMB::Dcerpc::Samr::SamprEncryptedUserPasswordNew.encrypt_password(
            computer_password,
            @simple.client.application_key
          )
        }
      )
    )
    samr_con[:samr].samr_set_information_user2(
      user_handle: user_handle,
      user_info: user_info
    )

    user_info = RubySMB::Dcerpc::Samr::SamprUserInfoBuffer.new(
      tag: RubySMB::Dcerpc::Samr::USER_CONTROL_INFORMATION,
      member: RubySMB::Dcerpc::Samr::UserControlInformation.new(
        user_account_control: RubySMB::Dcerpc::Samr::USER_WORKSTATION_TRUST_ACCOUNT
      )
    )
    samr_con.samr.samr_set_information_user2(
      user_handle: user_handle,
      user_info: user_info
    )
    print_good("Successfully created #{samr_con.domain_name}\\#{computer_name}")
    print_good("  Password: #{computer_password}")
    print_good("  SID:      #{get_computer_sid(samr_con, computer_name)}")
    report_creds(samr_con.domain_name, computer_name, computer_password)

    ComputerInfo.new(computer_name, computer_password)

  rescue RubySMB::Dcerpc::Error::SamrError => e
    raise MsSamrUnknownError, "A DCERPC SAMR error occured: #{e.message}"
  ensure
    if samr_con
      samr_con.samr.close_handle(user_handle) if user_handle
      samr_con.samr.close_handle(samr_con.domain_handle) if samr_con.domain_handle
      samr_con.samr.close_handle(samr_con.server_handle) if samr_con.server_handle
    end
  end

  def delete_computer(opts = {})
    tree = opts[:tree] || connect_ipc

    samr_con = connect_samr(tree)

    computer_name = opts[:computer_name] || datastore['COMPUTER_NAME']
    if computer_name.blank?
      raise MsSamrBadConfigError, 'Unable to delete the computer account since its name is unknown'
    end

    details = samr_con.samr.samr_lookup_names_in_domain(domain_handle: samr_con.domain_handle, names: [ computer_name ])
    raise MsSamrBadConfigError, 'The specified computer was not found.' if details.nil?
    details = details[computer_name]

    user_handle = samr_con.samr.samr_open_user(domain_handle: samr_con.domain_handle, user_id: details[:rid])
    samr_con.samr.samr_delete_user(user_handle: user_handle)
    print_good('The specified computer has been deleted.')
  rescue RubySMB::Dcerpc::Error::SamrError => e
    # `user_handle` only needs to be closed if an error occurs in `samr_delete_user`
    # If this method succeed, the server took care of closing the handle
    samr_con.samr.close_handle(user_handle) if user_handle
    raise MsSamrUnknownError, "Could not delete the computer #{computer_name}: #{e.message}"
  ensure
    samr_con.samr.close_handle(samr_con.domain_handle) if samr_con.domain_handle
    samr_con.samr.close_handle(samr_con.server_handle) if samr_con.server_handle
  end

  def lookup_computer(opts = {})
    tree = opts[:tree] || connect_ipc

    samr_con = connect_samr(tree)

    computer_name = opts[:computer_name] || datastore['COMPUTER_NAME']
    if computer_name.blank?
      raise MsSamrBadConfigError, 'Unable to lookup the computer account since its name is unknown'
    end

    sid = get_computer_sid(samr_con, computer_name)
    print_good("Found #{samr_con.domain_name}\\#{computer_name} (SID: #{sid})")
  ensure
    samr_con.samr.close_handle(samr_con.domain_handle) if samr_con.domain_handle
    samr_con.samr.close_handle(samr_con.server_handle) if samr_con.server_handle
  end


  module_function

  def connect_ipc
    begin
      connect
    rescue Rex::ConnectionError => e
      raise MsSamrConnectionError, e.message
    end

    begin
      smb_login
    rescue Rex::Proto::SMB::Exceptions::Error, RubySMB::Error::RubySMBError => e
      raise MsSamrAuthenticationError, "Unable to authenticate ([#{e.class}] #{e})."
    end
    report_service(
      host: rhost,
      port: rport,
      host_name: simple.client.default_name,
      proto: 'tcp',
      name: 'smb',
      info: "Module: #{fullname}, last negotiated version: SMBv#{simple.client.negotiated_smb_version} (dialect = #{simple.client.dialect})"
    )

    begin
      simple.client.tree_connect("\\\\#{sock.peerhost}\\IPC$")
    rescue RubySMB::Error::RubySMBError => e
      raise MsSamrConnectionError, "Unable to connect to the remote IPC$ share ([#{e.class}] #{e})."
    end
  end

  def connect_samr(tree)
    begin
      vprint_status('Connecting to Security Account Manager (SAM) Remote Protocol')
      samr = tree.open_file(filename: 'samr', write: true, read: true)

      vprint_status('Binding to \\samr...')
      samr.bind(endpoint: RubySMB::Dcerpc::Samr)
      vprint_good('Bound to \\samr')
      server_handle = samr.samr_connect
    rescue RubySMB::Dcerpc::Error::FaultError => e
      elog(e.message, error: e)
      raise MsSamrUnexpectedReplyError, "Connection failed (DCERPC fault: #{e.status_name})"
    end

    if datastore['SMBDomain'].blank? || datastore['SMBDomain'] == '.'
      all_domains = samr.samr_enumerate_domains_in_sam_server(server_handle: server_handle).map(&:to_s).map(&:encode)
      all_domains.delete('Builtin')
      if all_domains.empty?
        raise MsSamrNotFoundError, 'No domains were found on the SAM server.'
      elsif all_domains.length > 1
        print_status("Enumerated domains: #{all_domains.join(', ')}")
        raise MsSamrBadConfigError, 'The SAM server has more than one domain, the target must be specified.'
      end

      domain_name = all_domains.first
      print_status("Using automatically identified domain: #{domain_name}")
    else
      domain_name = datastore['SMBDomain']
    end

    domain_sid = samr.samr_lookup_domain(server_handle: server_handle, name: domain_name)
    domain_handle = samr.samr_open_domain(server_handle: server_handle, domain_id: domain_sid)

    SamrConnection.new(samr, server_handle, domain_handle, domain_name)

  rescue RubySMB::Dcerpc::Error::DcerpcError => e
    elog(e.message, error: e)
    raise MsSamrUnexpectedReplyError, e.message
  rescue RubySMB::Error::RubySMBError
    elog(e.message, error: e)
    raise MsSamrUnknownError, e.message
  end

  def random_hostname(prefix: 'DESKTOP')
    "#{prefix}-#{Rex::Text.rand_base(8, '', ('A'..'Z').to_a + ('0'..'9').to_a)}$"
  end

  def report_creds(domain, username, password)
    service_data = {
      address: datastore['RHOST'],
      port: datastore['RPORT'],
      service_name: 'smb',
      protocol: 'tcp',
      workspace_id: myworkspace_id
    }

    credential_data = {
      module_fullname: fullname,
      origin_type: :service,
      private_data: password,
      private_type: :password,
      username: username,
      realm_key: Metasploit::Model::Realm::Key::ACTIVE_DIRECTORY_DOMAIN,
      realm_value: domain
    }.merge(service_data)

    credential_core = create_credential(credential_data)

    login_data = {
      core: credential_core,
      status: Metasploit::Model::Login::Status::UNTRIED
    }.merge(service_data)

    create_credential_login(login_data)
  end

  def get_computer_sid(samr_con, computer_name)
    details = samr_con.samr.samr_lookup_names_in_domain(
      domain_handle: samr_con.domain_handle,
      names: [ computer_name ]
    )
    raise MsSamrNotFoundError, 'The computer was not found.' if details.nil?

    details = details[computer_name]
    samr_con.samr.samr_rid_to_sid(
      object_handle: samr_con.domain_handle,
      rid: details[:rid]
    ).to_s
  end

end
end
