A transparent proxy prototype using tproxy written in python3 & trio

TransparentProxy.py

#!/usr/bin/env python3

import sys
import traceback
import trio
import trio.socket as socket
import struct

"""

ip rule add fwmark 0x1 lookup 100
ip route add local default dev lo table 100

iptables -t mangle -A OUTPUT -p udp -m mark --mark 0x42 -j RETURN
iptables -t mangle -A OUTPUT -p udp -j MARK --set-mark 0x1

iptables -t mangle -A PREROUTING -p udp -m mark --mark 0x1 -j TPROXY --on-port=8881 --on-ip=127.0.0.1

"""

IP_TRANSPARENT = 19
IP_ORIGDSTADDR = 20
IP_RECVORIGDSTADDR = IP_ORIGDSTADDR
SO_MARK = 36

socket_table = {}

def get_dst(ancdata):
    dstip = None
    for cmsg_level, cmsg_type, cmsg_data in ancdata:
        if cmsg_level == socket.SOL_IP and cmsg_type == IP_ORIGDSTADDR:
                family, port = struct.unpack('=HH', cmsg_data[0:4])
                port = socket.htons(port)
                if family == socket.AF_INET:
                    start = 4
                    length = 4
                ip = socket.inet_ntop(family, cmsg_data[start:start + length])
                dstip = (ip, port)
                break
    return dstip

def creat_socket_tuple(srcip, dstip):
    remote = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    remote.setsockopt(socket.SOL_SOCKET, SO_MARK, 0x42)

    local = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    local.setsockopt(socket.IPPROTO_IP, IP_TRANSPARENT, 1)
    local.setsockopt(socket.SOL_SOCKET, SO_MARK, 0x42)

    ev = trio.Event()
    socket_table[(srcip, dstip)] = (local, remote, ev)

    return (local, remote, ev)


def get_socket_tuple(srcip, dstip):
    socket_tuple = socket_table.get((srcip, dstip))
    if socket_tuple:
        return (socket_tuple, False)

    return (creat_socket_tuple(srcip, dstip), True)

async def inbound_loop(local, remote, srcip, cancel_scope):
    while True:
        response, dstip = await remote.recvfrom(4096)
        print("[ ] Got response from remote udp://%s:%d" % (dstip[0], dstip[1]))
        await local.sendto(response, srcip)
        cancel_scope.deadline = trio.current_time() + 30

async def outbound_loop(event, cancel_scope):
    while True:
        await event.wait()
        event.clear();
        cancel_scope.deadline = trio.current_time() + 30

async def proxy_data(srcip, dstip, data):

        socket_tuple, new = get_socket_tuple(srcip, dstip)
        local, remote, ev = socket_tuple
        await remote.sendto(data, dstip)
        ev.set() #keep alive

        if not new:
            print("[ ] existing connection")
            return

        print("[ ] new connection")

        try:
            await local.bind(dstip)
            async with trio.open_nursery() as nursery:
                cancel_scope = nursery.cancel_scope
                nursery.start_soon(inbound_loop, local, remote, srcip, cancel_scope)
                nursery.start_soon(outbound_loop, ev, cancel_scope)

            if nursery.cancel_scope.cancelled_caught:
                #timeout
                print("[-] Timeout for udp://%s:%d" %(dstip[0], dstip[1]))

        except:
            print("[*] Unexpected error:", sys.exc_info()[0])
            traceback.print_exc()
        finally:
            local.close()
            remote.close()
            del socket_table[(srcip, dstip)]


async def main():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.setsockopt(socket.SOL_IP, IP_RECVORIGDSTADDR, 1)
    s.setsockopt(socket.IPPROTO_IP, IP_TRANSPARENT, 1)

    await s.bind(('127.0.0.1', 8881))
    print("[+] Bound to udp://127.0.0.1:8881")

    async with trio.open_nursery() as nursery:
        while True:
            data, ancdata, _, srcip = await s.recvmsg(4096, socket.CMSG_SPACE(24))
            dstip = get_dst(ancdata)

            if dstip:
                print("[ ] Connection from udp://%s:%d to udp://%s:%d" % (srcip[0], srcip[1], dstip[0], dstip[1]))
                nursery.start_soon(proxy_data, srcip, dstip, data)
            else:
                print("[*] Failed to get original dest address")

if __name__ == "__main__":
    trio.run(main)