boxtools/router.py

import sys
import socket
from select import select
from serial import serial_for_url
from subprocess import Popen, PIPE, check_output, call
from configparser import ConfigParser

FILTER = "iptables -i enp4s0 -p tcp -m tcp --dport %d -j ACCEPT"

BASIC = """
iptables -P INPUT %(accept)s
iptables -P FORWARD %(accept)s
iptables -P OUTPUT ACCEPT
iptables -A INPUT -i lo -j ACCEPT
"""

sim = False

def unix_cmd(command):
    if sim:
        print('> %r' % command)
    else:
        print('$ %r' % command)
        return Popen(command.split(), stdout=PIPE).communicate()[0].decode()


class IoHandler:
    client = None
    handler = None

    def __init__(self, client, handler):
        self.handler = handler
        self.client = client

    def request(self):
        try:
            data = self.client.recv(1024)
            if data:
                self.write(data)
                return
        except Exception as e:
            print('ERROR in request: %s' % e)
        self.close()
        self.handler.close_client(self)

    def reply(self):
        try:
            self.client.sendall(self.read())
            return
        except ConnectionResetError:
            pass
        except Exception as e:
            print('ERROR in reply: %s' % e)
        self.close()
        self.handler.close_client(self)

    
class TcpHandler(IoHandler):
    def __init__(self, client, handler):
        self.socket = socket.create_connection(handler.addr)
        self.fno = self.socket.fileno()
        super().__init__(client, handler)

    def read(self):
        data = self.socket.recv(1024)
        if not data:
            raise ConnectionResetError('disconnected')
        return data

    def write(self, data):
        self.socket.sendall(data)

    def close(self):
        try:
            self.socket.close()
        except Exception as e:
            print('ERROR in close: %s' % e)


class SerialHandler(IoHandler):
    def __init__(self, client, handler):
        self.serial = serial_for_url(handler.addr, timeout=10)
        self.serial.timeout = None
        self.fno = self.serial.fileno()
        super().__init__(client, handler)

    def read(self):
        return self.serial.read(self.serial.in_waiting)

    def write(self, data):
        self.serial.write(data)

    def close(self):
        self.serial.close()


class AcceptHandler:
    """handler for routing
    
    :param: port offered port for routing
    :param: addr where to route
    :param: iocls the io handler class, currently TcpHandler or SerialHandler
    :param: maxcount the maximal number of concurrent connections. defauls to 1
            as a side effect, if the destination is a web server, the traffic
            are serialized (only one connection at a time), which helps for
            some moxa device servers. might be a problem, if connections are
            reused: in this case maxcount has to be increased ...
    """
    readers = {}

    def __init__(self, port, addr, iocls, maxcount=1):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind(('0.0.0.0', port))
        s.listen()
        self.socket = s
        self.addr = addr
        self.iocls = iocls
        self.readers[s.fileno()] = self.accept
        self.port = port
        self.available = maxcount
        self.pending = 0
        print('listening at port %d for %s(%r)' % (port, iocls.__name__, addr) )

    def close_client(self, iohandler):
        self.readers.pop(iohandler.fno, None)
        try:
            client = iohandler.client
            self.readers.pop(client.fileno())
            client.close()
        except Exception as e:
            print('ERROR in close_client: %s' % e)
        iohandler.client = None
        iohandler.fno = None
        self.available += 1
        if self.pending:
            self.pending -= 1
            self.accept()

    def accept(self):
        if not self.available:
            self.pending += 1
            return
        try:
            client, addr = self.socket.accept()
            handler = self.iocls(client, self)
        except Exception as e:
            print('ERROR creating %s(%r)' % (self.iocls.__name__, self.addr))
            client.close()
            return
        self.readers[client.fileno()] = handler.request
        self.readers[handler.fno] = handler.reply
        self.available -= 1

    @classmethod
    def run(cls, routes, restrict=None):
        if restrict is not None:
            lines = BASIC % dict(accept='DROP' if restrict else 'ACCEPT')
            unix_cmd('iptables -F')
            for line in lines.split('\n'):
                if line.strip():
                    unix_cmd(line)
            if restrict:
                unix_cmd(FILTER % 22)

        for port, dest in routes.items():
            port=int(port)
            if restrict:
                unix_cmd(FILTER % port)
            if '/' in dest:
                AcceptHandler(port, dest, SerialHandler)
            else:
                host, _, remoteport = dest.partition(':')
                if remoteport:
                    remoteport = int(remoteport)
                else:
                    remoteport = port
                AcceptHandler(port, (host, remoteport), TcpHandler)
        while True:
            try:
                ready, _, _ = select(cls.readers, [], [])
            except Exception as e:
                for r in cls.readers:
                    try:
                        select([r], [], [], 0.1)
                    except Exception as e:
                        print(r, e)
                raise
            for fno in ready:
                cls.readers[fno]()



if __name__ == '__main__':
    parser = ConfigParser()
    parser.read('/root/aputools/servercfg/%s.cfg' % socket.gethostname())
    AcceptHandler.run(parser['ROUTER'])