#!/usr/bin/env python3

import sys
import pytun
import struct
import ipaddress
import json
import socket

MAGIC = b'\x16\xd9\x6b\x52'

class Config:
    def __init__(self, data):
        if data['mode'] == 'server':
            self.is_server = True
        elif data['mode'] == 'client':
            self.is_server = False
        else:
            raise RuntimeError('invalid mode "%s"' % data['mode'])

        self.address = data['address']
        self.port    = data['port']

        self.iface_name    = data['iface_name']
        self.iface_netmask = data['iface_netmask']
        self.iface_mtu     = data['iface_mtu']

        self.iface_addr    = ipaddress.IPv4Address(data['iface_addr'])
        self.iface_dstaddr = ipaddress.IPv4Address(data['iface_dstaddr'])

class IpPacket:
    def __init__(self, header, body):
        self.header = header
        self.body = body

    def to_byte_string(self):
        return self.header.to_byte_string() + self.body

class IpHeader:
    def __init__(self, raw):
        self.ver_ihl, \
            self.tos, \
            self.total_length, \
            self.ident, \
            self.flags_fragoffset, \
            self.ttl, \
            self.proto, \
            self.chksum, \
            src, \
            dst, \
            self.opt1, \
            self.opt2, \
            self.pad = struct.unpack('>BBHHHBBHIIHBB', raw)

        self.src = ipaddress.IPv4Address(src)
        self.dst = ipaddress.IPv4Address(dst)

    def to_byte_string(self):
        return struct.pack(
            '>BBHHHBBH4s4sHBB',
            self.ver_ihl,
            self.tos,
            self.total_length,
            self.ident,
            self.flags_fragoffset,
            self.ttl,
            self.proto,
            self.chksum,
            self.src.packed,
            self.dst.packed,
            self.opt1,
            self.opt2,
            self.pad,
        )

def main():
    config_filename = sys.argv[1]

    with open(config_filename) as config_file:
        config = Config(json.load(config_file))

    tun_iface = pytun.TunTapDevice()

    print(tun_iface.name)

    tun_iface.addr    = str(config.iface_addr)
    tun_iface.dstaddr = str(config.iface_dstaddr)
    tun_iface.netmask = config.iface_netmask
    tun_iface.mtu     = config.iface_mtu

    tun_iface.up()

    addr_and_port = (config.address, config.port)

    if config.is_server:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        print('binding to %s, port %s' % addr_and_port)
        sock.bind(addr_and_port)
        sock.listen(1)
        print('waiting for incoming connection')
        conn, client_addr = sock.accept()
        print('accepted incoming connection from %s' % client_addr[0])

    else:
        conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        print('connecting to %s, port %s' % addr_and_port)
        conn.connect(addr_and_port)
        print('connected!')

    while True:
        buf = tun_iface.read(tun_iface.mtu)

        flags = buf[:2]
        proto = buf[2:4]

        if proto != b'\x08\x00':
            continue

        ip_packet = IpPacket(IpHeader(buf[4:28]), buf[28:])

        handle_ip_packet(ip_packet, conn, tun_iface, config)

def handle_ip_packet(ip_packet, conn, tun_iface, config):
    print(
        ip_packet.header.src,
        '→',
        ip_packet.header.dst,
        'len: %d' % ip_packet.header.total_length,
    )

    ip_packet_packed = ip_packet.to_byte_string()
    ip_packet_length = struct.pack('>I', len(ip_packet_packed))

    buf = MAGIC + ip_packet_length + ip_packet_packed

    if ip_packet.header.dst == config.iface_addr:
        print('sending to tun iface')
        tun_iface.write(buf)
    elif ip_packet.header.dst == config.iface_dstaddr:
        print('sending to remote peer')
        conn.send(buf)
    else:
        print('unknown destination, doing nothing')

if __name__ == '__main__':
    main()