#! /usr/bin/env python
# Script that process stdout from pcf and creates IRI messages
#
# Copyright (C) 2014 Martin Holkovic
#
# 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 sys
import socket

import threading
import subprocess
import ctypes
import time
import json
import os
import inspect
import random
import signal
import string
import getopt
import base64

sys.path.append(os.path.join(os.path.dirname(__file__), '../../../'))

from pcapy import *
from thread import *
from modules.sockets.li2 import connectUnixSocket
from parser_functions import *
from parser_protocols import *
from parser_client import *
from subprocess import call

debugMode = None # urcuje ci sa budu vypisovat spravy IRI aj pri spojenie s IRI-Core
runningMode = "online" # moznost nastavit aby sa spravy neposilali na IRI-Core ale na stdout
iriSocket = None # socket pre komunikaciu s IRI-Core
acks = {} # zoznam cisiel ACK Num z TCP hlavicky, potvrdzujuce nepotvrdene e-mailove adresy
unconfirmedData = {} # zoznam nepotvrdenych e-mailovych adries v ramci SMTP spojenia
interface = "any" # defaultne rozhranie, ktore bude modul odpocuvat

db = {}

def set(index, name, value):
	global db
	if not index in db:
		db[index] = {}
	db[index][name] = value

def get(index):
	global db
	if index in db:
		return db[index]
	return None
	

def delete(index):
	global db
	if index in db:
		del db[index]
	

def unset(index, name):
	global db
	if index in db and name in db[index]:
		del db[index][name]

def must_be_root():
	if os.geteuid() != 0:
		sys.stderr.write("You must be root to run this script.\n")
		sys.exit(0)


def handler(signum, frame):
    parser_disconnect()
    sys.exit(0)


def lineno():
    return inspect.currentframe().f_back.f_lineno


# pomocna funkcia na vypisovanie debugovych hlasok
def debug(string):
    if debugMode != None:
        unixTime = str(time.time())
        parts = unixTime.split('.')
        if len(parts[1]) == 1:
            parts[1] = parts[1] + "0"
        resultTime = time.strftime("%d %b %Y %H:%M:%S", time.localtime()) + "." + parts[1]
        print >> sys.stderr, resultTime + ": " + string


# funkcia na odoslanie spravy na socket
def send_to_socket(iriType, text, save, aditionalData):
    if aditionalData == None: # nebol detegovany ziadny e-mail (MAIL ani RCPT)
        aditionalData = ""
    message = "('smtp', "+str(time_nowF())+", '"+iriType+"', '"+text+"', [('TCP', ('"+save["clientIp"]+"', "+str(save["clientPort"])+", '"+save["serverIp"]+"', "+str(save["serverPort"])+"))"+aditionalData+"])"
    
    if runningMode == "online": # pri rezime offline sa neposielaju na IRI-IIF spravy
        iriSocket.send(message + "\n")
        debug("IRI-IIF: " + message)
    else: # pri debug mode sa IRI spravy vzpisuju na stdout
		print message


def stop_same_proccess(): # zastavi proces s rovnakym nazvom (nie je nutne zistovat PID a kill-ovat cez shell)
    pid = shell("ps uax | grep -v 'sudo' | grep 'python smtp.py' | grep -v 'grep' | grep -v 'stop' | head --lines=1 | sed 's/^[^ ]\+\s\+\([0-9]\+\)\s.*$/\\1/'")
    if pid != "" and pid != str(os.getpid()):
        os.system("kill " + pid)


def print_help(): # vypise napovedu
        print "sudo python smtp.py -s debug"
        print "    -s    offline rezim bez komunikacie s IRI-IIF (standalone)"
        print "    debug    zapnutie debugovania"


# vypocet unix time aktualneho casu spolu s milisekundami
def time_nowF():
    unixTime = str(time.time())
    parts = unixTime.split('.')
    if len(parts[1]) == 1:
        parts[1] = parts[1] + "0"
    return float(str(int(time.mktime(time.localtime()))) + "." + parts[1])


def parse_parameters(argv):
    global runningMode, debugMode, interface
    try:
        opts, remainder = getopt.getopt(argv[1:], "i:hsov", ["interface=","help", "standalone", "offline","verbose"]) # spracovanie parametrov
    except getopt.GetoptError, err:
        sys.stderr.write(str(err) + "\n")
        sys.exit(2)
    
    for param, value in opts: # prejdenie parametrov
        if param in("-i, --interface"):
            interface = value
        
        elif param in("-s, --standalone"): # spustenie bez komunikacie s IRI
            runningMode = "offline"
        
        elif param in ("-o,--offline"): # to iste ako parameter --standalone
            runningMode = "offline"
        
        elif param in ("-h,--help"): # napoveda
            runningMode = "help"
        
        elif param in ("-v,--verbose"): # vypis sprav IRI aj pri spojeni s IRI-Core
            debugMode = True
    
    if "debug" in remainder:
        debugMode = True
    if "stop" in remainder:
        runningMode = "stop"


def confirm(index, number): # potvrdit parameter MAIL alebo RCPT na zaklade SEQ a ACK cisel
    global acks
    if unconfirmedData[index].has_key(number): # existuje parameter, ktory ocakava potvrdenie tymto cislom
        if acks[index].has_key(number): # duplikovany paket
            return
        else:
            if unconfirmedData[index][number][0] == "MAIL": # potvrdzuje sa parameter MAIL
                set(index, "MAIL", unconfirmedData[index][number][1])
            else: # potvrdzuje sa parameter RCPT
                confirmed = get(index)
                if confirmed.has_key("RCPT"): # jeden parameter RCPT uz je potvrdeny
                    rcpt = confirmed["RCPT"]
                    rcpt.append(unconfirmedData[index][number][1])
                else: # prvy potvrdeny parameter RCPT
                    rcpt = [unconfirmedData[index][number][1]]
                set(index, "RCPT", rcpt)
    acks[index][number] = True # potvrdenie prislo skorej ako paramater


def add_data(index, number, category, value): # pridanie e-mailu k nepotvrdenych datam
    global unconfirmedData
    if acks[index].has_key(number):
        if unconfirmedData[index].has_key(number): # duplicitna sprava
            return
        else:
            if category == "MAIL": # e-mail odosielatela
                set(index, "MAIL", value)
            else: # e-mail prijimatela
                confirmed = get(index)
                if confirmed.has_key("RCPT"):
                    rcpt = confirmed["RCPT"]
                    rcpt.append(value)
                else:
                    rcpt = [value]
                set(index, "RCPT", rcpt)
    unconfirmedData[index][number] = [category, value]


def get_mail(data): # z obsahu spravy sa zisti e-mailova adresa odosielatela
    begin = string.find(data, "<")
    end = string.find(data, ">")
    return data[begin + 1 : end]


def get_rcpt(data):  # z obsahu spravy sa zisti e-mailova adresa prijimatela
    if data.count("@") > 1:
        begin = string.find(data, ":")
        end = string.find(data, ">")
        return data[begin + 1 : end]
    else:
        begin = string.find(data, "<")
        end = string.find(data, ">")
        return data[begin + 1 : end]


def create_smtp_nid(save): # z e-mailovych adries ulozenych v ramci spojenia sa vytvori parameter IRI spravy
    result = "('SMTP'"
    if "MAIL" in save:
        result += ", ('E-mail address', '"+save["MAIL"]+"')"
    else:
        result += ", ('E-mail address', 'unknown')"
    result += ", ["
    if "RCPT" in save:
		for receiver in save["RCPT"]:
		    result += ",('E-mail address', '"+str(receiver)+"')"
    result += "]"
    if False:
        result += ", ('E-mail Message-ID', 'TODO')"
    else:
        result += ", ('E-mail Message-ID', 'unknown')"
    result += ")"
    
    return result


def base64_data(data):
    data = data[:-2] # posledne dva znaky su \r \n
    if len(data) % 3 != 0:
        for i in range(len(data)):
            ascii = ord(data[i])
            if ( ascii < "a" or ascii > "z" ) and ( ascii < "A" or ascii > "Z" ) and ascii != "+" and ascii != "/" and ascii != "=":
                return False
        return True
    else:
        return False


def clear_and_delete(index): # odstrani vsetky udaje k danemu spojeniu
    del acks[index]
    del unconfirmedData[index]
    delete(index)


def get_save(lowerLayers): # najde ulozene spojenie
    index = make_index(lowerLayers, False)
    save = get(index)
    if save == None: # skusim prehodit odosielatela a prijimatela
        index = make_index(lowerLayers, True)
        save = get(index)
    
    if save == None: # spojenie nie je ulozene, vratim pomocnu hodnotu pre zistenie stavu
        save = {"state": 0}
    
    return save, index


def calc_ack_num(seqnum, length): # vzpocet hodnoty ACK ktora potvrdi aktualnu spravu
    return (seqnum + length) % 4294967296
    

def make_index(lowerLayers, invert): # vytvori index do uloziska na zaklade TCP flow
    if invert:
        return lowerLayers["srcIp"] + ":" + str(lowerLayers["srcPort"]) + "-" + lowerLayers["dstIp"] + ":" + str(lowerLayers["dstPort"])
    else:
        return lowerLayers["dstIp"] + ":" + str(lowerLayers["dstPort"]) + "-" + lowerLayers["srcIp"] + ":" + str(lowerLayers["srcPort"])


# funkcia na spracovania kazdeho prijateho paketu
# lowerLayers - identifikatory hlaviciek nizsich sietovych vrstiev
# header - TCP hlavicka
# data - TCP data
def callback(lowerLayers, header, data):
    global acks, unconfirmedData
    
    data = hex_to_str(data) # data sa prenasaju v HEXa podobe
    save, index = get_save(lowerLayers) # zistim si status aktualneho toku
    
    fin = header["flags"] & 1 # FIN flag z TCP hlavicky
    
    
    if save["state"] == 0: # naviazane nove TCP spojenie
        syn = header["flags"] & 2 # SYN flag z TCP hlavicky
        if syn > 0:
            set(index, "state", 1)
    
    elif save["state"] == 1: # naviazane TCP spojenie
        syn = header["flags"] & 2
        ack = header["flags"] & 16
        acks[index] = {}
        unconfirmedData[index] = {}
        #print str_to_hex(data) + " " + data
        if data.startswith('220'):
            set(index, "serverIp", lowerLayers["srcIp"])
            set(index, "serverPort", lowerLayers["srcPort"])
            set(index, "clientIp", lowerLayers["dstIp"])
            set(index, "clientPort", lowerLayers["dstPort"])
            set(index, "state", 2)
        elif data.startswith('EHLO') or data.startswith('HELO'):
            set(index, "clientIp", lowerLayers["srcIp"])
            set(index, "clientPort", lowerLayers["srcPort"])
            set(index, "serverIp", lowerLayers["dstIp"])
            set(index, "serverPort", lowerLayers["dstPort"])
            set(index, "state", 3)
        elif syn + ack == 0:
            clear_and_delete(index)
    
    elif save["state"] == 2:
        syn = header["flags"] & 2
        ack = header["flags"] & 16
        if data.startswith('250'):
            send_to_socket("REPORT", "User has connected to the SMTP server", save, None)
            set(index, "state", 4)
        elif data.startswith('EHLO') or data.startswith('HELO'):
            set(index, "state", 3)
        elif syn + ack == 0:
            clear_and_delete(index)
    
    elif save["state"] == 3:
        if data.startswith('250'):
            send_to_socket("REPORT", "User has connected to the SMTP server", save, None)
            set(index, "state", 4)
        elif data.startswith('AUTH'):
            send_to_socket("REPORT", "User has connected to the SMTP server", save, None)
            set(index, "state", 5)
        elif data.startswith('MAIL') or data.startswith('Mail'):
            send_to_socket("REPORT", "User has connected to the SMTP server", save, None)
            add_data(index, calc_ack_num(header["seqNum"], len(data)), "MAIL", get_mail(data))
            set(index, "state", 8)
        elif data.startswith('5'):
            send_to_socket("REPORT", "The user unsuccessfully tried connect to server", save, None)
            set(index, "state", 2)
        elif fin > 0 or data.startswith('221') or data.startswith('QUIT'):
            send_to_socket("REPORT", "User has disconnected from the SMTP server", save, None)
            clear_and_delete(index)
        elif data.startswith('STARTTLS'):
            set(index, "state", 11)
    
    elif save["state"] == 4:
        if data.startswith('AUTH'):
            set(index, "state", 5)
        elif data.startswith('MAIL') or data.startswith('Mail'):
            add_data(index, calc_ack_num(header["seqNum"], len(data)), "MAIL", get_mail(data))
            set(index, "state", 8)
        elif data.startswith('250'):
            confirm(index, header["ackNum"])
            set(index, "state", 8)
        elif fin > 0 or data.startswith('221') or data.startswith('QUIT'):
            send_to_socket("REPORT", "User has disconnected from the SMTP server", save, None)
            clear_and_delete(index)
        elif data.startswith('STARTTLS'):
            set(index, "state", 11)
        elif data.startswith('334') or base64_data(data):
            set(index, "state", 6)
    
    elif save["state"] == 5:
        if fin > 0 or data.startswith('221') or data.startswith('QUIT'):
            send_to_socket("REPORT", "User has disconnected from the SMTP server", save, None)
            clear_and_delete(index)
        elif data.startswith('334') or base64_data(data):
            set(index, "state", 6)
    
    elif save["state"] == 6:
        if data.startswith('AUTH'):
            send_to_socket("REPORT", "The user was unsuccessful in authentication", save, None)
            set(index, "state", 5)
        elif data.startswith('MAIL') or data.startswith('Mail'):
            send_to_socket("BEGIN", "The user was successfully authenticated", save, None)
            add_data(index, calc_ack_num(header["seqNum"], len(data)), "MAIL", get_mail(data))
            set(index, "state", 8)
        elif data.startswith('235'):
            send_to_socket("BEGIN", "The user was successfully authenticated", save, None)
            set(index, "state", 8)
        elif data.startswith('535'):
            send_to_socket("REPORT", "The user was unsuccessful in authentication", save, None)
            set(index, "state", 7)
        elif fin > 0 or data.startswith('221') or data.startswith('QUIT'):
            send_to_socket("REPORT", "User has disconnected from the SMTP server", save, None)
            clear_and_delete(index)
    
    elif save["state"] == 7:
        if data.startswith('AUTH'):
            set(index, "state", 5)
        elif fin > 0 or data.startswith('221') or data.startswith('QUIT'):
            send_to_socket("REPORT", "User has disconnected from the SMTP server", save, None)
            clear_and_delete(index)
        elif data.startswith('334') or base64_data(data):
            set(index, "state", 6)
    
    elif save["state"] == 8:
        if data.startswith('250'):
            confirm(index, header["ackNum"])
        elif data.startswith('MAIL') or data.startswith('Mail'):
            add_data(index, calc_ack_num(header["seqNum"], len(data)), "MAIL", get_mail(data))
        elif data.startswith('RCPT') or data.startswith('Rcpt'):
            add_data(index, calc_ack_num(header["seqNum"], len(data)), "RCPT", get_rcpt(data))
            confirm(index, header["seqNum"])
        elif data.startswith('EHLO') or data.startswith('HELO'):
            unset(index, "MAIL")
            unset(index, "RCPT")
            set(index, "state", 3)
        elif data.startswith('DATA'):
            confirm(index, header["seqNum"])
            set(index, "state", 9)
        elif data.startswith('354'):
            set(index, "state", 9)
        elif data.startswith('RSET'):
            unset(index, "MAIL")
            unset(index, "RCPT")
            set(index, "state", 8)
        elif fin > 0 or data.startswith('221') or data.startswith('QUIT'):
            send_to_socket("END", "User has disconnected from the SMTP server", save, None)
            clear_and_delete(index)
    
    elif save["state"] == 9:
        if data.startswith('EHLO') or data.startswith('HELO'):
            unset(index, "MAIL")
            unset(index, "RCPT")
            set(index, "state", 3)
        elif data.startswith('RSET'):
            unset(index, "MAIL")
            unset(index, "RCPT")
            set(index, "state", 8)
        elif data.startswith('250'):
            send_to_socket("CONTINUE", "An e-mail of length TODO bytes successfully sent", save, create_smtp_nid(save))
            set(index, "state", 10)
        elif data.startswith('4') or data.startswith('5'):
            send_to_socket("CONTINUE", "The user unsuccessfully tried send a e-mail", save, create_smtp_nid(save))
            set(index, "state", 10)
        elif fin > 0 or data.startswith('221') or data.startswith('QUIT'):
            send_to_socket("END", "User has disconnected from the SMTP server", save, None)
            clear_and_delete(index)
    
    elif save["state"] == 10:
        if data.startswith('EHLO') or data.startswith('HELO'):
            unset(index, "MAIL")
            unset(index, "RCPT")
            set(index, "state", 3)
        elif data.startswith('MAIL') or data.startswith('Mail'):
            unset(index, "MAIL")
            unset(index, "RCPT")
            add_data(index, calc_ack_num(header["seqNum"], len(data)), "MAIL", get_mail(data))
            set(index, "state", 8)
        elif data.startswith('RSET'):
            unset(index, "MAIL")
            unset(index, "RCPT")
            set(index, "state", 8)
        elif fin > 0 or data.startswith('221') or data.startswith('QUIT'):
            send_to_socket("END", "User has disconnected from the SMTP server", save, None)
            clear_and_delete(index)
    
    elif save["state"] == 11:
        if data.startswith('220') or ord(data[0]) < 25: # v pripade zacati TLS je na prvej pozicii asci hodnota 20-24 reprezentujuca typ TLS obsahu
            send_to_socket("CONTINUE", "SSL connection initiated", save, None)
            set(index, "state", 10)
        elif data.startswith('4') or data.startswith('5'):
            set(index, "state", 4)
        elif data.startswith('AUTH'):
            set(index, "state", 5)
    
    #print db


if __name__ == "__main__":
    must_be_root()
    signal.signal(signal.SIGTERM, handler)
    parse_parameters(sys.argv)
    
    
    if runningMode == "help":
        print_help()
    elif runningMode == "stop": # ukonci proces s rovnakym nazvom
        stop_same_proccess()
    else:
        stop_same_proccess() # naraz moze byt spusteny len jeden proces
        if runningMode == "online": # v pripade online modu sa spoji s IRI
            iriSocket = connectUnixSocket("/tmp/iricol")
        parser_connect(interface, "TCP", sys.argv) # spojenie s parserom
        parser_read(callback) # citanie dat a volanie vlastnej funkcie na ne
        parser_disconnect()

