twisted tcp proxy

tcp-proxy.py

#!/usr/bin/env python
#Simple http header steganography project
import base64
from queue import Queue
from twisted.internet import defer, protocol, reactor

class ProxyClientProtocol(protocol.Protocol):
    def connectionMade(self):
        print("Client: connected to peer")
        self.cli_queue = self.factory.cli_queue
        self.cli_queue.get().addCallback(self.serverDataReceived)

    def serverDataReceived(self, chunk):
        if chunk is False:
            self.cli_queue = None
            print("Client: disconnecting from peer")
            self.factory.continueTrying = False
            self.transport.loseConnection()
        elif self.cli_queue:
            print("Client: writing %d bytes to peer" % len(chunk))
            self.transport.write(chunk)
            self.cli_queue.get().addCallback(self.serverDataReceived)
        else:
            self.factory.cli_queue.put(chunk)

    def dataReceived(self, chunk):
        print("Client: %d bytes received from peer" % len(chunk))
        self.factory.srv_queue.put(chunk)

    def connectionLost(self, why):
        if self.cli_queue:
            self.cli_queue = None
            print("Client: peer disconnected unexpectedly")

class ProxyClientFactory(protocol.ReconnectingClientFactory):
    maxDelay = 10
    continueTrying = False
    protocol = ProxyClientProtocol

    def __init__(self, srv_queue, cli_queue):
        self.srv_queue = srv_queue
        self.cli_queue = cli_queue

class ProxyServer(protocol.Protocol):
    queue = 0

    def steg(self, chunk):
        if(self.queue.empty()): return chunk
        response = chunk.decode("utf-8") 
        p = response.find('\r\n\r\n')
        if p >= 0:
            header = response[:p]
            header += "\r\nContent-MD5: "
            raw_bytes = bytes(self.queue.get(),'ISO-8859-1')
            base64_bytes = base64.b64encode(raw_bytes)
            base64_message = base64_bytes.decode('ascii')
            header += base64_message
            print(base64_message)
            res = header + response[p:]
            return res.encode('utf-8')
        return chunk

    def connectionMade(self):
        self.srv_queue = defer.DeferredQueue()
        self.cli_queue = defer.DeferredQueue()
        self.srv_queue.get().addCallback(self.clientDataReceived)
        factory = ProxyClientFactory(self.srv_queue, self.cli_queue)
        reactor.connectTCP("127.0.0.1", 8000, factory)

    def clientDataReceived(self, chunk):
        print("Server: writing %d bytes to original client" % len(chunk))
        self.transport.write(self.steg(chunk))
        self.srv_queue.get().addCallback(self.clientDataReceived)

    def dataReceived(self, chunk):
        print("Server: %d bytes received" % len(chunk))
        self.cli_queue.put(chunk)

    def connectionLost(self, why):
        self.cli_queue.put(False)

class ProxyFactory(protocol.Factory):
    def __init__(self, q):
        self.protocol = ProxyServer
        self.protocol.queue = q

def getSecretText():
    with open('Sofokles-Antygona.txt', 'r', encoding="ISO-8859-1") as file:
        queue = Queue()
        while True:
            chunk = file.read(16)
            if not chunk:
                break
            queue.put(chunk)
        return queue

if __name__ == "__main__":
    print("Proxy: initialized")
    q = getSecretText()
    factory = ProxyFactory(q)
    reactor.listenTCP(80, factory, interface="0.0.0.0")
    reactor.run()

Prosty projekt na przedmiot na pw implementujący steganografie sieciową, w nagłówku http,
co można poprawić: - jest jedna kolejka dla wszystkich - parametryzacja