require "mysql"
require 'win32ole'
# Unbrowse SNMP Scripting Interface Demo 
# by Vipin.K.Narayanan
#
# t2sql -    start, stop, and control the unbrowse trap 
#            passive receiver. stores all incoming traps 
#	     to a SQL database 
#
#
# Licensing
#   You can use this in any way you please. No warranties.
#   (c) Unleash Networks 2009, All rights reserved
# ---------------------------------------------------------------
if ARGV.size < 4
	puts "USAGE: ruby #{$0} <trapdatabase> <mysql-server-address> <mysql-user> <mysql-password>"
	puts "for empty password give \"\""
else 
	# ---------------------------------------------------
	# 
	#	Dumps almost all trap information onto screen
	# ---------------------------------------------------
	def pr_trap(one_trap)
		# --------------------------
		# Print all trap information 
		# --------------------------
		puts " ------------------------------------------" 
		puts " Trap ID    : #{one_trap.ID}"
		puts " From Agent : #{one_trap.AgentAddress}"
		puts " To Manager : #{one_trap.DestinationAddress}"
		puts " Timestamp  : #{one_trap.TimestampLocal}"
		puts " User/Comm  : #{one_trap.UserCommunity}"
		puts " Varbinds   : #{one_trap.VarbindCount}"
		puts " OID        : #{one_trap.EffectiveTrapOID}"


		# --------------------------
		# Print all varbinds in trap
		# --------------------------
		puts " --------- Varbind list ----------" 
		(0 .. one_trap.VarbindCount - 1).each do |i|
			one_varbind = one_trap.GetVarbindByIdx(i)
			puts "\t #{one_varbind.OID}   =  #{one_varbind.Value}"
		end
		puts " ------------------------------------------" 

	end

	# -------------------------------------------------------------
	# Connect to a mysql server 
	#  
	# -------------------------------------------------------------

	begin
     		# connect to the MySQL server
     		dbh = Mysql.real_connect(ARGV[1],ARGV[2],ARGV[3],ARGV[0])
     		# get server version string and display it
     		puts "Server version: " + dbh.get_server_info
	rescue Mysql::Error => e
     		puts "Error code: #{e.errno}"
     		puts "Error message: #{e.error}"
     		puts "Error SQLSTATE: #{e.sqlstate}" if e.respond_to?("sqlstate")
		exit
	end

	# ---------------------------------------------------
	# StoreTrapToSQL  :Store the incoming trap to a SQL Connection
	# Parameters
	#	thetrap - the trap that needs to be stored
	#	Sqldb_handle - the sql connection
	# ---------------------------------------------------

	def storetraptosql(sqldb_handle, thetrap)

		# -----------------------------
		# Create list of valid columns 
		# (see schema above)
		# -----------------------------
		colnames = "Timestamp,AgentAddress,EffectiveTrapOID,Version,Context,UserCommunity,VarbindCount"

		(0 .. thetrap.VarbindCount - 1).each do |j|
				colnames="#{colnames}" + ",Varbind_OID_#{j}, Varbind_Value_#{j}" 
		end

		# ---------------------------------------------------
		# Assign values from the incoming trap to the columns
		# ---------------------------------------------------
		colvalues = " \"#{thetrap.TimestampLocal}\", \"#{thetrap.AgentAddress}\" , \"#{thetrap.EffectiveTrapOID}\" , \"#{thetrap.SNMPVersion}\" , \"#{thetrap.ContextName}\" ,\"#{thetrap.UserCommunity}\" , \"#{thetrap.VarbindCount}\"" 
		
		(0 .. thetrap.VarbindCount - 1).each do |i|
			onevar = thetrap.GetVarbindByIdx(i)
			colvalues = colvalues +", \"#{onevar.OID}\" , \"#{onevar.Value}\""  
	
		end
	
		# ---------------------------------------------------
		# The complete INSERT statement
		# ---------------------------------------------------
		insertstmt = "INSERT INTO alltraps ( #{colnames} ) VALUES ( #{colvalues} )" 
		puts " SQL Statement = #{insertstmt}"	 

		# ---------------------------------------------------
		# Execute the INSERT statement
		# ---------------------------------------------------
		sqldb_handle.query("#{insertstmt}")
	end

	# --------------------------------------------------
	# Check usage & arguments
	# --------------------------------------------------
	MaxTraps   = 100000
	stopping   = false

	rep_mgr = WIN32OLE.new("UnbrowseSNMP.RepositoryManager")
	rep_db  = rep_mgr.LoadRepositoryReadOnly


	# --------------------------------------------------
	# Create the trap server  and attach the events 
	# --------------------------------------------------
	trap_mgr = WIN32OLE.new("UnbrowseSNMP.TrapReceiver")
	trap_mgrEvents = WIN32OLE_EVENT.new(trap_mgr,"_IVWSnTrapMgrEvents")
	puts "Loaded the Unbrowse SNMP Trap Server" 


	# --------------------------------------------------
	# Use the UDP Server mode (other options are Winpcap and Raw Sockets)
	# --------------------------------------------------
	trap_mgr.UDPServerMode = true


	# --------------------------------------------------
	# Open a new database - you can save this later
	# --------------------------------------------------
	trap_mgr.NewDatabase 

	# --------------------------------------------------
	# Start the passive trap receiver
	# --------------------------------------------------
	puts "Starting .." 
	stopping="false"
	trap_mgr.Start 
	puts "Listening for traps .." 

	# ------------------------------------------------------
	# Enter into an loop, processing traps
	# Here were stopping after a few traps for demo purposes
	# ------------------------------------------------------
	puts "Dumping traps as they are received .. (Ctrl+C) to quit" 

	trap_mgrEvents.on_event("TNF_NewTrap") do  |id,one_trap|
		pr_trap(one_trap)
		storetraptosql(dbh, one_trap)
	end

	# --------------------------------------------------
	# Go into event loop
	# --------------------------------------------------

	while true
    		WIN32OLE_EVENT.message_loop
    		sleep(0.1)
	end


	# --------------------------------------------------
	# Stop the trap receiver 
	# --------------------------------------------------
	puts "stopping the trap receiver"

	trap_mgr.Stop


	# --------------------------------------------------
	# Close the DB connection
	# --------------------------------------------------
	dbh.close if dbh

	puts "Done"
end
exit