leolee192
2/18/2016 - 11:23 AM

tcp relay

tcp relay

#!/usr/bin/python

from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
import logging


class Handler(BaseHTTPRequestHandler):
    def _set_headers(self):
        self.send_response(200)
        self.end_headers()

    def do_GET(self):
        logging.info('do_GET')
        self._set_headers()
        self.wfile.write('reply GET\n')

    def do_POST(self):
        logging.info('do_POST')
        content_len = int(self.headers.getheader('content-length', 0))
        post_body = self.rfile.read(content_len)
        self._set_headers()
        self.wfile.write('reply POST\n' + post_body + '\n')


def run(port):
    HTTPServer(('0.0.0.0', port), Handler).serve_forever()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
    import sys

    port = int(sys.argv[1])
    run(port)
#!/usr/bin/python

import logging
import socket
from threading import Thread


class PipeThread(Thread):
    pipes = []

    def __init__(self, source, sink):
        Thread.__init__(self)
        self.source = source
        self.sink = sink

        logging.info('new pipe thread %s (%s -> %s)' % \
                     (self, source.getpeername(), sink.getpeername()))
        PipeThread.pipes.append(self)
        logging.info('%s pipes active' % len(PipeThread.pipes))

    def run(self):
        while True:
            try:
                data = self.source.recv(1024)
                if not data:
                    break
                self.sink.send(data)
            except Exception as e:
                logging.error(e)
                break
        logging.info('%s terminating' % self)
        self.sink.shutdown(socket.SHUT_WR)
        PipeThread.pipes.remove(self)
        logging.info('%s pipes active' % len(PipeThread.pipes))


class Relay(Thread):
    def __init__(self, port, dst_host, dst_port):
        Thread.__init__(self)
        logging.info('relaying: localhost:%s -> %s:%s' % (port, dst_host, dst_port))
        self.dst_host = dst_host
        self.dst_port = dst_port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind(('0.0.0.0', port))
        self.sock.listen(1)

    def run(self):
        while True:
            newsock, address = self.sock.accept()
            logging.info('new session for %s %s' % address)
            fwd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            fwd.connect((self.dst_host, self.dst_port))
            PipeThread(newsock, fwd).start()
            PipeThread(fwd, newsock).start()


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
    logging.info('starting relay')

    import sys

    port = int(sys.argv[1])
    dst_host = sys.argv[2]
    dst_port = int(sys.argv[3])
    Relay(port, dst_host, dst_port).start()