from drawbridge import DrawBridge
from scapy.layers.http import HTTPResponse, HTTP
from scapy.all import IP, TCP
from scapy.packet import Raw
import json


def modify_websocket_chat(raw_packet):
    pkt = IP(raw_packet)
    
    tcp0 = pkt.getlayer(TCP)
    try:
        jsonb = bytes(tcp0.payload)[2:]
        if jsonb == b'':
            return raw_packet
    except IndexError:
        return raw_packet
    try:
        json.loads(jsonb.decode("utf-8"))
    except (json.decoder.JSONDecodeError, UnicodeDecodeError):
        return raw_packet
    
    tcp0.payload = Raw(bytes(tcp0.payload)[:2] + b'{"sender": "hackerboy23", "message": "boom"}')
    del pkt[IP].len
    del pkt[IP].chksum
    del pkt[TCP].chksum

    return bytes(pkt)


black_png = (
    b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\n\x00\x00'
    b'\x00\n\x08\x06\x00\x00\x00\x8d2\xcf\xbd\x00\x00\x00\x0cIDATx'
    b'\xda\xed\xc1\x01\r\x00\x00\x00\xc2\xa0\xf5H\xfd\x00\x00\x00'
    b'\x00IEND\xaeB`\x82'
)
black_payload = HTTP(bytes(
                "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\n\r\n", 'utf-8') + black_png
            )
black_len = str(len(black_png)).encode()



def modify_img_request(raw_packet):
    pkt = IP(raw_packet)

    if pkt.haslayer(HTTPResponse):
        http_layer = pkt.getlayer(HTTP)
        http_response = pkt.getlayer(HTTPResponse)

        if http_response.fields.get('Content_Type') == b'image/png':
            print("Found PNG file...")

            # http_layer.payload = black_payload
            pkt.show2()

            if pkt.haslayer(TCP):
                http_layer.fields['Content_Length'] = black_len
                del pkt[TCP].chksum
                del pkt[IP].chksum
                del pkt[IP].len
    else:
        return raw_packet
    return bytes(IP(bytes(pkt)))


db = DrawBridge()
db.add_queue(modify_websocket_chat, queue=2, src_port=80)
db.run()