#!/usr/bin/env python3

import time
from lib.tuntap import TunThread
from lib import dns
from lib import query
from lib.packet import Packet, PacketPool

hostname = b'vpn.bgpat.net'
query.Field.HostName.default = hostname
Packet.hostname = hostname

txpool = PacketPool()
rxpool = PacketPool()


class VPNServer(TunThread):
    daemon = True
    name = 'tun0'
    addr = '192.168.200.1'
    gateway = '192.168.200.1'

    def receive(self, data):
        global rxpool
        pkt = Packet(data)
        rxpool.put(pkt)


tun = VPNServer()
tun.start()


class DNSServer(dns.ServerThread):
    daemon = True

    def receive(self, req, addr, port, data):
        txInit = query.TxInitialize(req)
        txSend = query.TxSend(req)
        rxRecv = query.Receive(req)
        rxPoll = query.Polling(req)
        if txInit is not None:
            res = self.txinit(txInit['id'], txInit['count'])
        elif txSend is not None:
            res = self.txsend(txSend['id'], txSend['sequence'], txSend['data'])
        elif rxRecv is not None:
            res = self.rxrecv(rxRecv['id'], rxRecv['sequence'])
        elif rxPoll is not None:
            res = self.rxpoll()
        else:
            res = query.Error()
        return {
            'value': bytes(res),
            'type': res.type,
        }

    def txinit(self, id, count):
        global txpool
        txpool[id] = Packet(count)
        return query.Ok(count=count, sequence=count)

    def txsend(self, id, seq, data):
        global txpool
        pkt = txpool[id]
        pkt[seq] = data
        remain = pkt.count - len(pkt)
        if not remain:
            tun.send(pkt.unpack())
            del txpool[id]
        return query.Ok(count=remain, sequence=seq)

    def rxrecv(self, id, seq):
        global rxpool
        pkt = rxpool[id]
        if seq in pkt:
            del pkt[seq]
        keys = list(pkt.keys())
        if len(keys):
            i = keys[0]
            return query.RxSend(data=pkt[i], sequence=i, id=id)
        else:
            del rxpool[id]
            return query.Error()

    def rxpoll(self):
        global rxpool
        if rxpool.empty():
            return query.Error()
        pkt = rxpool.front()
        if pkt.id in rxpool:
            return query.Error()
        rxpool[pkt.id] = pkt
        rxpool.pop()
        return query.RxInitialize(data=pkt[0], count=pkt.count, id=pkt.id)


dnsd = DNSServer()
dnsd.start()

while True:
    time.sleep(1)

@@ -812,16 +812,23 @@
 f txInit
+.params
  is not 
@@ -825,32 +825,67 @@
 ms is not None:%0A
+            params = txInit.params%0A
             res 
@@ -902,22 +902,22 @@
 nit(
-txInit
+params
 %5B'id'%5D, 
 txIn
@@ -912,22 +912,22 @@
 %5B'id'%5D, 
-txInit
+params
 %5B'count'
@@ -948,16 +948,23 @@
 f txSend
+.params
  is not 
@@ -961,32 +961,67 @@
 ms is not None:%0A
+            params = txSend.params%0A
             res 
@@ -1038,22 +1038,22 @@
 end(
-txSend
+params
 %5B'id'%5D, 
 txSe
@@ -1048,22 +1048,22 @@
 %5B'id'%5D, 
-txSend
+params
 %5B'sequen
@@ -1068,22 +1068,22 @@
 ence'%5D, 
-txSend
+params
 %5B'data'%5D
@@ -1103,16 +1103,23 @@
 f rxRecv
+.params
  is not 
@@ -1116,32 +1116,67 @@
 ms is not None:%0A
+            params = rxRecv.params%0A
             res 
@@ -1193,22 +1193,22 @@
 ecv(
-rxRecv
+params
 %5B'id'%5D, 
 rxRe
@@ -1203,22 +1203,22 @@
 %5B'id'%5D, 
-rxRecv
+params
 %5B'sequen
@@ -1242,16 +1242,23 @@
 f rxPoll
+.params
  is not 

Fix query switch