campus-total-block/dst-domain-match-check.rb

#!/usr/bin/env ruby

require 'inifile'
require 'pp'

require "resolv"
require "mysql2"

require "socket"
require "thread"
require "syslog"
require "rubygems"
require "open-uri"
require "ipaddress"

STDOUT.sync = true
Syslog.open("#{$PROGRAM_NAME}", Syslog::LOG_PID)

def log(msg)
  Syslog.log(Syslog::LOG_ERR, "%s", msg)
  STDERR.puts("STDERR: [ #{msg} ]") if $debug
end


$my_dir = __dir__

config_filename = ARGV[0]

if config_filename.nil? or !File.exists?(config_filename)
	$debug = true
	log("The config file: [\"#{config_filename}\"] doesn't exist")
	exit 1
end

config_file = IniFile.load(config_filename)

$value_location = config_file["main"]["value"].to_i

$table_name = config_file["mysql"]["table"]

$debug = false

if config_file["main"]["debug"].to_s == "1"
	$debug = true
end

log("Started with DEBUG => #{$debug}")

socket_type = config_file["socket"]["type"]
socket_address = config_file["socket"]["address"]
port = config_file["socket"]["port"] if config_file["socket"]["port"]

$mysql_hostname = config_file["mysql"]["host"]
$mysql_username = config_file["mysql"]["username"]
$mysql_password = config_file["mysql"]["password"]
$mysql_db = config_file["mysql"]["database"]

trap "SIGINT" do
  STDERR.puts "STDERR: Exiting"
  exit 130
end


def cleanupUnixSocket(unix_socket_address)
  if File.exist?(unix_socket_address)
    File.delete(unix_socket_address)
  end
end

def generate_domain_array(domain)
  parts = domain.split(".")
  l = parts.length

  domains = []
  domains << parts.join(".")

  i = 0
  while i < l do
        domains << ".#{parts[i..-1].join(".")}"
    i = i +1
  end
  domains
end

def check_domain_in_db(domain)
  domains =  generate_domain_array(domain)
  begin
    db_client = Mysql2::Client.new(:host => $mysql_hostname, :username => $mysql_username, :password => $mysql_password, :database => $mysql_db)
  rescue => e
    log(e)
    log(e.inspect)
  end
  sql_query = "SELECT * FROM #{$table_name} WHERE dstdom IN ( #{domains.map { |d| "'#{d}'" }.join(" , ")} )"
  log("Running SQL query: [ \"#{sql_query}\" ]")
  results = db_client.query(sql_query)
  if results.nil?
	db_client.close
	return false
  end
  if results.size > 0
    db_client.close
    return true
  end
  db_client.close
  return false
end

def check_domain(domain)
  log("checking domain #{domain}") if $debug
  db_res = check_domain_in_db(domain)
  log("DB RES #{db_res}") if $debug

  if db_res
    return true
  end

  return false
end

def requestTest(request)
  return if request == nil
  request = request.split
  matched_to = []
  ret = "0"
  log("Request size: #{request.size} , value_location #{$value_location}") if $debug
  if request.size >= $value_location
    begin
      # check if the requst is a domain name or ip address
      is_ip_address = IPAddress.valid?(request[$value_location])
      if !is_ip_address
        if check_domain(request[$value_location].downcase)
          ret = "1"
        end
      end
    rescue => e
      log(e)
      log(e.inspect)
      ret = "1"
    end
  end
  return { "request_id" => request[0], "ret" => "#{ret}" }
end


def validr?(request)
  if request.ascii_only? && request.valid_encoding?
    return true
  else
    STDERR.puts("errorness line [ #{request} ]")
    return false
  end
end

answers = { "0" => "ERR", "1" => "OK" }


log("Socket binding is starting")
case socket_type
when /^tcp/i
  begin
    log("Trying to bind: #{socket_address}:#{port}")
    server_socket = TCPServer.new(socket_address, port)

    loop do
      Thread.start(server_socket.accept) do |s|
        log("#{s} is accepted")
        processingtQueue = Queue.new

        proccessor = Thread.new do
          loop do
            incomming_request = processingtQueue.pop
            return if incomming_request.nil?
            Thread.new do
              result = requestTest(incomming_request) if validr?(incomming_request)
              s.puts("#{result["request_id"]} #{answers[result["ret"]]}")
              log("result for request: #{s} => [ #{incomming_request} ] , res => #{result}") if $debug
            end
          end
        end

        while line = s.gets
          processingtQueue << line.strip.chomp
          log("original request: #{s} => [ #{line.chomp} ]") if $debug
        end
        proccessor.join
        log("#{s} is gone")
        s.close
      end
    end
  rescue => e
    log(e)
    log(e.inspect)
    exit 10
  end
when /^unix/i
  begin
    log("starting to bind UNIX Socket")
    if IPAddress.valid?(socket_address)
      log("Cannot use IP address #{socket_address} for unix socket")
      exit 1
    end
	address_url = URI.parse(socket_address)
	unix_socket_address = address_url.path

    cleanupUnixSocket(unix_socket_address)
    log("Trying to bind unix socket")
    server_socket = UNIXServer.new(unix_socket_address)
    log("UNIX Socket was bounded")

    loop do
      Thread.start(server_socket.accept) do |s|
        log("#{s} is accepted")
        while line = s.gets
          line = line.strip.chomp
          log("original request: #{s} => [ #{line} ]") if $debug
          result = requestTest(line) if validr?(line)
          s.puts("#{result["request_id"]} #{answers[result["ret"]]}")
          log("result for request: #{s} => [ #{line} ] , res => #{result}") if $debug
        end
        log("#{s} is gone")
        s.close
      end
    end
  rescue => e
    File.delete(unix_socket_address) if File.exists?(unix_socket_address)
    log(e)
    log(e.inspect)
    exit 11
  end
else
  log("Sokcet type: #{socket_type}  is not supported")
  exit 1
end