#!/usr/local/bin/ruby

=begin
-----------------------------------------------------------------------------
"THE BEER-WARE LICENSE":

<johnny@netvor.sk> wrote this file. As long as you retain this notice you
can do whatever you want with this stuff. If we meet some day, and you think
this stuff is worth it, you can buy me a beer (*or Kofola) in return.   

  johnny ^_^ <johnny@netvor.sk>
-----------------------------------------------------------------------------

sshguard2.rb
============

 - more dynamic (or drastic?) ruby equivalent of sshguard
 - improved heuristics
 - only ipfw support right now, feel free to expand

install
-------
1. copy somewhere
2. make executable
3. add:
     auth.info;authpriv.info     | exec /path/to/this/file
   to your syslog.conf & reload

config
------
put values in $Conf under /usr/local/etc/sshguard2.conf 
(or override $ConfDefaults in this file)

looking for updates? 
--------------------
http://netvor.sk/~johnny/

looking for help?
-----------------
http://www.google.com/
johnny@netvor.sk

have fun ^^/

=end

require 'syslog'
require 'thread'
require 'socket'

# some constants
CONF_FILE = '/usr/local/etc/sshguard2.conf'

# Firewall interface base class
class Firewall 
	# block access 
	def block ip
		Syslog::info "blocking access from #{ip}"
	end
	# allow access
	def release ip
		Syslog::info "release #{ip}"
	end
	# is access blocked
	def blocked? ip
		return false
	end
	# array of blocked addresses
	def blockedIPs
		return []
	end
end

class IpfwBlacklistFirewall < Firewall
	# block ip address
	def block ip
		return if $fw.blocked? ip

		`ipfw table 22 add #{ip}`
		print "block"
		cmd = "echo '#{ip}' >> '"+$Conf['blacklist_file']+"'"
		`#{cmd}`
	end

	# not in my blacklist :)
	def release ip
	end

	# is given ip addr blocked?
	def blocked? ip
		cmd = "grep '#{ip}' '"+$Conf['blacklist_file']+"'"
		`#{cmd}`
		return ($?.to_i == 0)
	end
end

# ipfw interface
class IpfwFirewall < Firewall
	def initialize
		@rule_map = {}

		# parse existing firewall rules and fill in @rule_map
		# TODO:
		`ipfw show`.each { |line|
			if line =~ /^([0-9]+) .* deny ip from ([0-9\.]+) to any 22/
				if $1.to_i.between? $Conf['first_rule'],$Conf['last_rule']
					@rule_map[$2.strip] = $1.to_i
				end
			end
		}
	end

	# deny access from given ip address
	def block ip
		# don't block already blocked ip
		return if $fw.blocked? ip

		# unique rule_id and ip->rule mapping
		rule = rand($Conf['last_rule'] - $Conf['first_rule']) + $Conf['first_rule']
		@rule_map[ip] = rule;

		# block access to ip on firewall
		`ipfw add #{rule} deny all from #{ip} to any 22`

		super
	end

	# allow access from given ip address
	def release ip
		rule = @rule_map[ip]
		@rule_map.delete_if { |k,v| v == rule }
	
		# release it
		`ipfw delete #{rule}`

		super
	end

	# is given ip addr blocked?
	def blocked? ip
		return ! @rule_map[ip].nil?
	end

	# get list of blocked hosts
	def blockedIPs
		return @rule_map.keys
	end
end

class Host 
	def initialize ip
		@ip = ip
		@fail = []
		@success = []
		@block_until = Time.now - 1
		@last_seen = Time.now 
	end	

	# failed login
	def fail
		@last_seen = Time.now
		@fail << Time.now
	end

	# successful login
	def success
		@success << Time.now
	end

	# should we release this host?
	def release? 
		return false unless $fw.blocked? @ip
		return Time.now > @block_until
	end

	# decide if we should block access from us or not
	def block?
		# maul attack -- block
		block = true if 
			@fail.size >= $Conf['block_count']

		# hit-and-run attack (distributed dictionary attacks) -- block
		block = true if 
			@success.size == 0 and
			@fail.size >= $Conf['delayed_block_cont'] and
		   	Time.now > @last_seen + $Conf['delayed_block_after'] 

		# don't block already blocked
		block = false if $fw.blocked? @ip

		# block this address?
		if block
			@block_until = Time.now + $Conf['release_after'] + $Conf['release_after_exp']*@fail.size
			Syslog::notice "block #{@ip} until #{@block_until}"
		end

		return block
	end

	# forget about this host
	def forget
		@fail.delete_if { |t| Time.now > t + $Conf['forget_fail_after'] }
		@success.delete_if { |t| Time.now > t + $Conf['forget_success_after'] }
	end
end

# is given ip whitelisted?
def whitelist? ip
	if not defined? $Conf['whitelist'][ip]
		return false
	end
	Syslog::info "ignore #{ip} -- whitelisted"
	return true
end

# make sure we die properly
Thread.abort_on_exception = true

# hosts list and mutex for insertion
$hosts = {}
$mutex = Mutex.new

# seed
srand

# default config, don't modify. use Conf in /usr/local/etc/sshguard2.conf
$ConfDefault = {
	'firewall'			=> 'ipfw_blacklist',
	'first_rule'			=> 1000,
	'last_rule'			=> 1999,
	'block_count'			=> 6,
	'delayed_block_cont'		=> 1,
	'delayed_block_after'		=> 300,
	'release_after'			=> 600,
	'release_after_exp'		=> 3,
	'forget_fail_after'		=> 3600*24,
	'forget_success_after' 		=> 3600*24*7,
	'blacklist_file'		=> '/crypto/jail/www/www/public_html/johnny/sshbl.krabica',
}	

# load config
$Conf = $ConfDefault
load CONF_FILE if FileTest.readable? CONF_FILE

# verify and parse config
case $Conf['firewall']
when "ipfw"
	$fw = IpfwFirewall.new
when "ipfw_blacklist"
	$fw = IpfwBlacklistFirewall.new
else
	raise "unknown firewall type #{$Conf['firewall']}"
end

# TODO: validate $Conf and rewrite invalid vals with those in $ConfDefault

# open syslog
Syslog::open 'sshguard2', Syslog::LOG_PID, Syslog::LOG_AUTH
Syslog::notice 'process started'

# load all 'already blocked' hosts
$fw.blockedIPs.each { |ip|
	$hosts[ip] = Host.new ip
}	

Thread.new {
	loop do
		sleep(($Conf['release_after'] / 3) % 3600)  # at least once in a hour
		$hosts.each { |ip,host|
			# forgive host previous login errors? maybe...
			host.forget

			# allow access if host is blocked and should be released
			$fw.release ip if host.release?
		}
	end
}

# read STDIN
unless STDIN.tty?
	while((line = STDIN.gets))
		# parse this line
		case line.strip
		when /error: PAM: authentication error for (\w+) from (.*)$/,
		     /error: PAM: authentication error for illegal user (\w+) from (.*)$/,
			/Invalid user (\w+) from (.*)$/
			ip, ok = $2, false
		when /Accepted publickey for (\w+) from (.*) port ([0-9]*) ssh2$/,
			 /Accepted keyboard-interactive\/pam for (\w+) from (.*) port ([0-9]*) ssh2$/
			ip, ok = $2, true
		else
			ip = nil
		end

		# skip unknown lines
		next if ip.nil?

		# lookup host addr
		ip = IPSocket.getaddress(ip).strip

		# whitelist?
		next if whitelist? ip

		# lookup host
		if(host = $hosts[ip]).nil?
			host = Host.new ip
			$hosts[ip] = host 
		end

		if ok
			host.success
		else
			host.fail
			$fw.block ip if host.block?
		end
		
		# give some time to others 
		Thread.pass
	end
else
	puts 'STDIN wtf??'
end 

# cleanup
Syslog::info 'i don\'t blame you'
Syslog::close