#!/usr/bin/env python3
#vtli.py a script for calculating the threat level on a scale of 1-8 and sending it to the arduino
#Joe McManus josephmc@cmu.edu
#version 0.2 2012/02/15
#Copyright (C) 2012 Joe McManus
#This program is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with this program.  If not, see <http://www.gnu.org/licenses/>.

import socket
from datetime import datetime, timedelta
import io
import sys
import platform
import os
import re

host = "10.0.0.177"
port = 1563
snortLog="alert.csv"

maxAlerts = 34 #What should be considered the highest threat level, turn for your environ

if platform.python_version() < "3.0.0": 
        printUsage("Python 3.0 or greater is required for this to run. Sorry")

def printUsage(error):
	print("ERROR: " + error)
	print("USAGE: " + sys.argv[0])
	sys.exit()

def putScale(host, port, scale):
	socket.setdefaulttimeout(10)
	s= sockets = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	try:
		s.connect((host, port))
	except: 
		printUsage("Unable to connect to Arduino")
	try: 
		#t = s.send(str.encode(str(scale)))
		t = s.send(str.encode(str(scale)))
	except: 
		printUsage("Unable to send data to Arduino")
	s.close()

def checkFile(snortLog): 
	if os.path.isfile(snortLog):
		print("OK: Found Snort Log- " + snortLog); 
	else:
		printUsage("Log File Not Found");

def createRegEx(): #We generate a regex for the last 24 hours.
	today=datetime.now()
	expression=str(today.strftime("%m/%d-%H"))
	i=0
	for i in range(23):
		newTime=datetime.now() - timedelta(hours=i)
		expression=expression + '|' + str(newTime.strftime("%m/%d-%H"))
	return expression

def countAlerts(snortLog):
	alertCount=0
	today=datetime.now()
	regex=re.compile(('^(%s)') % createRegEx()) 
	try:	
		fh = open(snortLog, 'r')
	except: 
		printUsage("Unable to open Snort Log")

	for line in fh.readlines():
		if regex.search(line):
			alertCount=alertCount+1
	return alertCount

def getLevel(maxAlerts, alertCount):
	score=round(alertCount/maxAlerts*9)
	if score > 9:
		score = 9 #Can only display a max of 10 leds, but we start counting at 0.
	if score < 1:
		score =1  #We can never have 0 threats...
	return score
		
checkFile(snortLog)	
alertCount=countAlerts(snortLog)
print("OK: Found " + str(alertCount) + " alerts in log file" )
level=getLevel(maxAlerts,alertCount)
print("-Connecting to display at: " + host)
putScale(host, port, level)
print("OK: Update Completed")