# frozen_string_literal: true

require 'socket'
require 'stringio'

DNS_TYPES = {
  1 => "A",
  2 => "NS",
  5 => "CNAME",
}

class DNSEncapsulationOfWriteOperations
  # fixed format, just hardcode
  private def make_question_header(query_id)
    [query_id, 0x0100, 0x0001, 0x0000, 0x0000, 0x0000].pack('nnnnnn')
  end

  public def make_dns_query(domain, query_id)
    question = domain.split('.').map { |label| [label.length, label].pack('Ca*') }.join + "\x00"
    question << [0x0001, 0x0001].pack('nn') # query type and class
    make_question_header(query_id) + question
  end
end

class DNSEncapsulationOfReadOperations
  def read_domain_name(buf)
    domain = []

    loop do
      label_length = buf.read(1).unpack1('C')
      break if label_length == 0

      if (label_length & 0xc0) == 0xc0
        # DNS compression
        pointer_offset = ((label_length & 0x3f) << 8) + buf.read(1).unpack1('C')
        old_pos = buf.pos
        buf.pos = pointer_offset
        domain << read_domain_name(buf)
        buf.pos = old_pos
        break
      else
        # Normal label
        domain << buf.read(label_length)
      end
    end

    domain.join('.')
  end

  def read_rdata(buf, length)
    case DNS_TYPES[@type] || @type
    when "CNAME", "NS"
      read_domain_name(buf)
    when "A"
      buf.read(4).unpack1('C4')
    else
      buf.read(length)
    end
  end

end

class DNSHeader
  attr_reader :id, :flags, :num_questions, :num_answers, :num_auth, :num_additional

  def initialize(buf)
    hdr = buf.read(12)
    @id, @flags, @num_questions, @num_answers, @num_auth, @num_additional = hdr.unpack('nnnnnn')
  end
end

class DNSRecord
  attr_reader :name, :type, :class, :ttl, :rdlength, :rdata, :parsed_rdata

  def initialize(buf)
    @dns_read_ops = DNSEncapsulationOfReadOperations.new
    @name = @dns_read_ops.read_domain_name(buf)
    @type, @class, @ttl, @rdlength = buf.read(10).unpack('nnNn')
    @parsed_rdata = read_rdata(buf, @rdlength)
  end

  def read_rdata(buf, length)
    @type_name = DNS_TYPES[@type] || @type
    if @type_name == "CNAME" or @type_name == "NS"
      @dns_read_ops.read_domain_name(buf)
    elsif @type_name == "A"
      buf.read(length).unpack('C*').join('.')
    else
      buf.read(length)
    end
  end

  def to_s
    "#{@name}\t\t#{@ttl}\t#{@type_name}\t#{@parsed_rdata}"
  end
end

class DNSQuery
  attr_reader :domain, :type, :cls

  def initialize(buf)
    dns_read_ops = DNSEncapsulationOfReadOperations.new
    @domain = dns_read_ops.read_domain_name(buf)
    @type, @cls = buf.read(4).unpack('nn')
  end
end

class DNSResponse
  attr_reader :header, :queries, :answers, :authorities, :additionals

  def initialize(bytes)
    buf = StringIO.new(bytes)  # 将字节数组打包成字符串并创建 StringIO 对象
    @header = DNSHeader.new(buf)
    @queries = (1..@header.num_questions).map { DNSQuery.new(buf) } if @header.num_questions > 0
    @answers = (1..@header.num_answers).map { DNSRecord.new(buf) } if @header.num_answers > 0
    @authorities = (1..@header.num_auth).map { DNSRecord.new(buf) } if @header.num_auth > 0
    @additionals = (1..@header.num_additional).map { DNSRecord.new(buf) } if @header.num_additional > 0
  end
end

def send_dns_query(sock, domain, query_id)
  dns_write_ops = DNSEncapsulationOfWriteOperations.new
  sock.setsockopt(Socket::SOL_SOCKET, Socket::SO_RCVTIMEO, [3, 0].pack('l_2')) # timeout - 3s
  sock.send(dns_write_ops.make_dns_query(domain, query_id), 0)
end

def receive_dns_response(sock)
  reply, _ = sock.recvfrom(1024)
  DNSResponse.new(reply)
end

module Resolvfast
  class Error < StandardError; end

  def self.main(domain)
    begin
      sock = UDPSocket.new
      sock.bind('0.0.0.0', rand(10240..65535))
      sock.connect('223.5.5.5', 53)

      # send query
      send_dns_query(sock, domain, 1)

      # receive & parse response
      response = receive_dns_response(sock)
      print "Answers for #{domain}:\n"
      response.answers.each do |record|
        puts record
      end
    ensure
      sock.close
    end
  end
end

if ARGV.empty?
  puts "Usage: #{$PROGRAM_NAME} <domain>"
  exit 1
end

Resolvfast.main(ARGV[0])

$ ruby resolvfast.rb google.com
Answers for google.com:
google.com              143     A       172.217.160.110
$ ruby resolvfast.rb baidu.com
Answers for baidu.com:
baidu.com               212     A       39.156.66.10
baidu.com               212     A       110.242.68.66