# -*- coding: utf8 -*-

# Copyright (C) Christian Barral López - All Rights Reserved
# Unauthorized copyright of this file, via any medium is strictly prohibited
# Proprietary and confidential
# Written by Christian Barral <cbarrallopez@gmail.com>, April 2019

import lib.helpers.directory_helper as directory_helper

from lib.helpers.colors import *
from lib.sniffer.NetworkProfile import NetworkProfile
from lib.sniffer.ARPPoisoner import ARPPoisoner

from scapy.all import *


class Sniffer:
    """A network sniffer implementation using scapy's built-in functions. This class
    executes two important functionalities. First of all, it sniffs the hole network thanks to
    the ARP Poisoning performed by the lib.sniffer.ARPPoisoner class, applying several options
    gathered by the lib.core.ConfigurationLoader class and passed on the __init__ function, and exporting
    the packets to a .pcap file when the sniffing finishes. If an input_file is specified (which
    has to be a .csv generated by the lib.Scanner.Scanner class) the sniffer will also process each
    packet captured to see if it fits with the information gathered on that file thanks to the
    Network Profile class.

        :param interface: Host network interface to use for sniffing
        :param gateway: Gateway's IP address. Used by ARPPoisoner
        :param targets: Target's to sniff from. Used by ARPPoisoner
        :param localnet: Specify localnet sniffing. Used by ARPPoisoner
        :param randomize: Target a number of random hosts for sniffing. Used by ARPPoisoner.
        :param arp_type: Specify the type of ARP Packets to use. Used by ARPPoisoner.
        :param output_directory: Directory where the output file is going to be written.
        :param output_file: Output file name
        :param sniff_input_file: .csv file to contrast packets captured.
        :param bpf_filter: Berkeley Packet Filter to apply when sniffing.
        :param timeout: Time to spend sniffing.
        :param packet_count: Number of packets to sniff. 0 means infinite.
        :param verb: Verbose level (0-3). 0 means no output.
        :type interface: str
        :type gateway: str
        :type targets: list
        :type localnet: bool
        :type randomize: int
        :type arp_type: int, str
        :type output_directory: str
        :type output_file: str
        :type sniff_input_file: str
        :type bpf_filter: str
        :type timeout: int
        :type packet_count: int
        :type verb: int

    Apart from class __init__ arguments, there are another attributes.
    Attributes:
        __ARPPoisoner: lib.sniffer.ARPPoisoner. Instance used to poison network hosts.
        __network_profile: lib.sniffer.NetworkProfile. Object representing the expected behaviour from the network.

    see_also::
        lib.sniffer.ARPPoisoner for ARP Poisoning
        lib.scanner.Scanner to know about scan CSV file
        lib.core.ConfigurationLoader to see the application arguments
        lib.sniffer.NetworkProfile to know how the network profile is built.

    warning::
        The BPF Filter must follow some explicit rules. Check biot.com/capstats/bpf.html to
        see all the examples.
    """

    def __init__(self,
                 interface,
                 gateway,
                 targets,
                 localnet,
                 randomize,
                 arp_type,
                 output_directory,
                 output_file,
                 sniff_input_file,
                 bpf_filter,
                 timeout,
                 packet_count,
                 verb):

        self.__timeout = timeout
        self.__packet_count = packet_count
        self.__verb = verb
        self.__filter = bpf_filter
        self.__output_directory = directory_helper.test_directory(output_directory)
        self.__output_file = output_file
        self.__arp_poisoner = ARPPoisoner(interface, gateway, targets, localnet,
                                          randomize, arp_type, verb)

        # If sniff input file used, create the network profile
        if sniff_input_file is not None:
            self.__network_profile = NetworkProfile(sniff_input_file, self.__output_directory)
        else:
            self.__network_profile = None

        # Scapy configuration
        conf.verb = verb
        conf.iface = interface

    @property
    def timeout(self):
        return self.__timeout

    @property
    def packet_count(self):
        return self.__packet_count

    @property
    def verb(self):
        return self.__verb

    @property
    def filter(self):
        return self.__filter

    @property
    def output_directory(self):
        return self.__output_directory

    @property
    def output_file(self):
        return self.__output_file

    @property
    def arp_poisoner(self):
        return self.__arp_poisoner

    def __build_file_name(self):
        """Builds the output file depending on the output directory and desired file name, appending the current
        date time and ending the file with the '.pcap' extension if it does not have it.
        """
        file_base_name = self.output_file
        current_datetime = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

        # Remove extension, if it exists
        if '.' in file_base_name:
            file_base_name = re.sub('\..*', '', self.output_file)

        return directory_helper.process_directory(self.output_directory) + '/' + file_base_name + '_' + current_datetime + '.pcap'

    def __print_sniffer_config(self):
        """Prints sniffing configurations depending on the instance attributes.
        """
        if self.timeout is not None:
            print_warning('Sniffing timeout: {}' .format(self.timeout))

        if self.packet_count != 0:
            print_warning('Packet count: {}'.format(self.packet_count))
        else:
            print_warning('Packet count: Infinite')

        if self.filter:
            print_warning('Packet filter: {}'.format(self.filter))

        if self.__network_profile:
            print_warning('Built network profile from input file. Log will be generated')

        print_warning('Verbose level: {}'.format(self.verb))

    def start_sniffing(self):
        """Initializes the sniffing action itself.
        """

        # Variable to store the packets
        sniffed_packets = None
        # Output file
        output = self.__build_file_name()
        # Print the configuration to terminal
        self.__print_sniffer_config()
        # Stablish callback function if there is a network profile
        if self.__network_profile is not None:
            callback = self.__network_profile.analyze_packet
        else:
            callback = None
        print_success('Writing output to {}'.format(output))
        try:
            self.arp_poisoner.run_poisoner()
            sniffed_packets = sniff(prn=callback,
                                    filter=self.filter,
                                    iface=conf.iface,
                                    count=self.packet_count,
                                    timeout=self.timeout)
        except KeyboardInterrupt:
            print_fail("User stopped sniffing. Captured packets written to {}" .format(output))
        except Exception as e:
            halt_fail('{}'.format(e))
        finally:
            if sniffed_packets is not None:
                try:
                    wrpcap(output, sniffed_packets)
                except Exception as e:
                    print_fail('Error while exporting packets {}'.format(e))

            # On stop, internally call the restore_network() method
            self.arp_poisoner.stop = True
            self.arp_poisoner.restore_network()
