dnet-whois/lib/registry.rb

# Class, which provides possibilities to act directly with the data in the registry.
class Registry
  require 'logger'
  require 'ipaddr'

  MAX_LIMIT = 40_000
  IP_PATHS = %w[route route6 inetnum inet6num].freeze
  SEARCH_PATHS = %w[as-set aut-num dns key-cert mntner organisation person registry role route-set schema tinc-key].freeze
  DIRECT_NAME = 'direct'.freeze

  def initialize(args)
    # Ensure data types of arguments
    raise ArgumentError, 'Wrong argument format' unless args.is_a? Hash
    raise ArgumentError, 'Directory must be a string' unless args[:dir].is_a? String
    raise ArgumentError, 'Logger must be a logger' unless args[:logger].is_a? Logger

    # Process and save the arguments
    @data_dir = "#{args[:dir]}#{File::SEPARATOR}data#{File::SEPARATOR}"
    @logger = args[:logger]

    @registry_data = {}

    # Check if the registry is available
    raise ArgumentError, 'Path does not exist' unless File.exist? @data_dir
  end

  # Reads the registry again
  def update
    limit_counter = 0
    registry_data = {}

    Dir["#{@data_dir}#{File::SEPARATOR}**#{File::SEPARATOR}*"].each do |file|
      if File.file? file
        if File.readable? file
          # Read file and save the lines into a array
          cnt = File.readlines file

          # Remove line endings
          cnt.map!(&:chomp)

          # Remove empty lines
          cnt.reject!(&:empty?)

          # Concat lines with CRLF
          val = cnt.join "\r\n"

          # Determine the path (something like aut-num or inetnum)
          splited_dir = File.dirname(file).split(File::SEPARATOR)
          if splited_dir.empty?
            @logger.error "Invalid path for file #{file}"
            next
          elsif splited_dir[-1] == 'data'
            # file is direct in data path, e. g. filter files
            path = DIRECT_NAME
          else
            path = splited_dir[-1]
          end

          # Determine the object something like AS1 or ::/0
          object_name = File.basename file

          # Create the path in the registry data if needed
          registry_data[path] = {} unless registry_data.key? path

          if IP_PATHS.include? path
            # Extract the IP out of the filename
            ip_address_str = object_name.gsub('_', '/')
            begin
              # Parse the IP
              ip_address = IPAddr.new ip_address_str
              # Save the parsed ip and the file content
              registry_data[path][object_name] = [ip_address, val]
            rescue IPAddr::InvalidAddressError
              @logger.error "Failed to parse #{ip_address_str}"
            end
          else
            # Save the file content
            registry_data[path][object_name] = val
          end
        else
          @logger.error "File #{file} is not readable."
        end
      end

      # Avoid crashing the program by adding
      # a lot of objects to the registry
      limit_counter += 1
      if limit_counter > MAX_LIMIT
        @logger.error 'Reached file limit.'
        break
      end
    end

    @logger.info 'Save new registry data.'
    @registry_data = registry_data
  end

  # Searches for an IP address in the registry
  #
  # @return [Array]
  def lookup_ipaddr(query)
    raise SuspiciousBehavior, 'Tried to lookup a IP address which is not a IP address' unless query.is_a? IPAddr

    if query.ipv4?
      route_path = 'route'
      inetnum_path = 'inetnum'
    elsif query.ipv6?
      route_path = 'route6'
      inetnum_path = 'inet6num'
    else
      raise SuspiciousBehavior, "Valid IP address is not IPv4 or IPv6: #{query}"
    end

    results = []

    results << lookup_ip_in_path(query, route_path)
    results[-1].map! do |value|
      [value[0], 2, value[1], value[2]]
    end

    results << lookup_ip_in_path(query, inetnum_path)
    results[-1].map! do |value|
      [value[0], 1, value[1], value[2]]
    end

    results.flatten!(1)

    results.reject!(&:empty?)

    results.sort! do |ip_a, ip_b|
      (ip_a[0] == ip_b[0] ? ip_a[1] <=> ip_b[1] : ip_a[0] <=> ip_b[0])
    end

    # Remove sorting information
    results.map! do |val|
      [val[2], val[3]]
    end

    return results
  end

  # Searches for an IP address in a specific path in the registry
  #
  # @return [Array]
  def lookup_ip_in_path(query, path)
    raise ArgumentError, 'Invalid IP path' unless IP_PATHS.include? path
    raise ArgumentError, 'IP address must be of type IP address' unless query.is_a? IPAddr

    results = []

    @registry_data[path].each do |key, value|
      ipaddr = value[0]
      data = value[1]

      if ipaddr.include? query
        # Found match
        results << [ipaddr, "#{path}/#{key}", data]
      end
    end

    return results
  end

  # Searches for files which are directly in the data directory.
  #
  # @return [Array]
  def lookup_direct(query)
    raise SuspiciousBehavior, 'Tried to lookup a string which is not a string' unless query.is_a? String

    query = query.sanitize.downcase
    return lookup_in_path(query, DIRECT_NAME)
  end

  # Searches for an object in a specific path in the registry
  #
  # @return [Array]
  def lookup_in_path(query, path)
    raise ArgumentError, 'Query must be a string' unless query.is_a? String
    raise ArgumentError, 'Path must be a string' unless path.is_a? String

    results = []

    @registry_data[path].each do |key, data|
      if key.downcase == query
        # Match found
        results << ["#{path}/#{key}", data]
      end
    end

    return results
  end

  # Searches for an object in the registry
  #
  # @return [Array]
  def lookup(query)
    raise SuspiciousBehavior, 'Tried to lookup a string which is not a string' unless query.is_a? String

    query = query.sanitize.downcase
    results = []

    SEARCH_PATHS.each do |search_path|
      results << lookup_in_path(query, search_path)
    end

    results.flatten!(1)

    return results
  end
end