#! /usr/bin/env python
# Class describing interception transferred through HI1
#
# Copyright (C) 2011 Matej Gregr, Michal Kajan, Libor Polcak,
#                    Vladimir Vesely, Petr Kramolis
#
# 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 2
# of the License.
# 
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import getopt
import os
import socket
import sys
import time
from threading import *

import scapy.all as scapy

from modules.iriiif.dhcp_status import *
from modules.sockets.li2 import connectUnixSocket

request_list = {}
lease_list = {}
clean_time = 600
check_timeout = 5
last_check = 0
lock = Lock()
lease_multiplier = 1 

IRICOL = "/tmp/iricol"

Types = {
    "discover":1,
    "offer":2,
    "request":3,
    "decline":4,
    "ack":5,
    "nak":6,
    "release":7,
    "inform":8,
    "force_renew":9,
    "lease_query":10,
    "lease_unassigned":11,
    "lease_unknown":12,
    "lease_active":13,
 }

def sendReport(report):
    global s
    i = 0;
    while(i<6):
        try:
            s.send(report + "\n")

        except socket.error, e:
            s = connectUnixSocket(IRICOL)
            i+=1;
            if(i>0):
                print "DHCPv4 module: trying to reconect to IRI-CORE"
                time.sleep(1);
            continue
        break

def array2strMac(array):
    mac = ""
    mac += "%.2x" % ord(array[0])
    for char in array[1:6]:
        mac += ":%.2x" % ord(char)
    return mac

class Timer(Thread):
    def run(self):
        while(1):
            time.sleep(check_timeout)
            with lock:
                tmp = []
                #check of lease
                for mac in lease_list:
                    lease = lease_list[mac]
                    if lease[0] < time.time():
                        #send end msg
                        dhcp_info = DHCPStatus("END",mac,lease[1],time.time(),"IP address assignment time-outed")
                        sendReport(repr(dhcp_info))
                        tmp.append(mac)

                for mac in tmp:
                    del lease_list[mac]

                tmp = []
                #clean request_list
                for xid in request_list:
                    req=request_list[xid]
                    if req[1] < time.time():
                       tmp.append(xid)

                for xid in tmp:
                    del request_list[xid]


def check_addr(src_mac, addr):
    global lease_list
    for src in lease_list:
        lease=lease_list[src]
        if lease[1]==addr:
            if src == src_mac:
               #lease time extended - send continue
                dhcp_info = DHCPStatus("CONTINUE",src_mac,addr,time.time(),"client renewed IP address assignment")
                sendReport(repr(dhcp_info))
                return
        else:
            if src == src_mac:
               #there is old lease for this addres - send end for it
                dhcp_info = DHCPStatus("END",src,addr,time.time(),"IP address reassigned to another client")
                sendReport(repr(dhcp_info))
    
    dhcp_info = DHCPStatus("BEGIN",src_mac,addr,time.time(),"new IP address assigned")
    sendReport(repr(dhcp_info))
        

def scapy_callback(scapy_packet):
    global last_check
    global request_list
    global lease_list

    dhcp=scapy_packet.getlayer(scapy.DHCP,1)
    if dhcp is  None:
        return

    #ether=scapy_packet.getlayer(Ether,1)
    bootp=scapy_packet.getlayer(scapy.BOOTP,1)
    mac = array2strMac(bootp.chaddr)

    #replace mac with client id option if it exists
    options = dhcp.options
    for x in options:
        if x[0] == 'client_id':
            mac = "0x" + x[1].encode("hex")
    addr = bootp.yiaddr
    #scapy_packet.show()

    for msg_typ in options:
        if msg_typ[0] == 'message-type':
            with lock:
                if msg_typ[1]==Types["discover"]:
                    dhcp_info = DHCPStatus("REPORT",mac, 0 ,time.time(), "client looked for server (DHCP discover)")
                    sendReport(repr(dhcp_info))
        
                elif msg_typ[1]==Types["request"]:
                    request_list[bootp.xid] = (mac,time.time()+clean_time)
                    dhcp_info = DHCPStatus("REPORT",mac, 0 ,time.time(), "client requested IP address assignment")
        
                elif msg_typ[1]==Types["ack"]:
        
                    #get client mac addres
                    if bootp.xid in request_list.keys():
                        src = request_list[bootp.xid]
                        src = src[0]
                        check_addr(src, addr)
                        del request_list[bootp.xid]
        
                        #get lease_time
                        for x in options:
                            if x[0] == 'lease_time':
                                lease_list[src] = (time.time() + (x[1] * lease_multiplier),addr)
        
                elif msg_typ[1]==Types["nak"]:
                    if bootp.xid in request_list.keys():
                        src = request_list[bootp.xid]
                        src = src[0]
                        del request_list[bootp.xid]
                        dhcp_info = DHCPStatus("REPORT", src, 0 ,time.time(),"IP address assignment refused")
                        sendReport(repr(dhcp_info))
        
                elif msg_typ[1]==Types["release"]:
                    dhcp_info = DHCPStatus("END", mac, 0 ,time.time(),"IP address released")
                    sendReport(repr(dhcp_info))

def help():
    print "-h \t\t print this help"
    print "-i <interface>\t sniff on given interface"
    print "-p <file>\t read pcap file"


if __name__ == "__main__":
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hi:p:", ["help", "interface=","pcap="])
    except getopt.GetoptError, err:
        sys.stderr.write(str(err) + "\n")
        help()
        sys.exit(2)

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


    t = Timer()
    last_check = time.time()
    t.start()
    #open socket
    #time.sleep(3)
    global s
    s = connectUnixSocket(IRICOL)

    for o, a in opts:
        if o in("-i, --interface"):
            #use given interface
            sys.stderr.write("Using interface: %s\n" % a)
            if a == "any":
                scapy.sniff(prn=scapy_callback,store=0, filter="port 68 or port 67 and udp")
            else:
                scapy.sniff(prn=scapy_callback, iface=a ,store=0, filter="port 68 or port 67 and udp")
            sys.exit()

        elif o in ("-h", "--help"):
            help()
            sys.exit()

        elif o in ("-p,--pcap"):
            #read pcap
            sys.stderr.write("Reading pcap file: %s\n" % a)
            scapy.sniff(prn=scapy_callback, offline=a ,store=0, filter="port 68 or port 67 and udp")
            sys.exit()

    #capture packets on all interfaces
    scapy.sniff(prn=scapy_callback,store=0, filter="port 68 or port 67 and udp")
