import sys

from twisted.internet.protocol import Protocol
from twisted.internet.ssl import CertificateOptions, PrivateCertificate
from twisted.protocols.loopback import loopbackTCP
from twisted.internet.task import react
from twisted.internet import reactor

class Client(Protocol):
    def connectionMade(self):
        self.count = 1000
        self.transport.startTLS(CertificateOptions(verify=False, requireCertificate=False))
        assert sys.argv[1] in ("small", "big")
        self.next_iteration()

    def next_iteration(self):
        self.count -= 1
        if sys.argv[1] == "small":
            for i in range(1_000):
                self.transport.write(b"X")
        else:
            self.transport.write(b"X" * 1000)

        if self.count == 0:
            self.transport.loseConnection()
        else:
            reactor.callLater(0, self.next_iteration)

class Server(Protocol):
    received = 0

    def connectionMade(self):
        with open("server.pem", "rb") as f:
            options = PrivateCertificate.loadPEM(f.read()).options()
        self.transport.startTLS(options)

    def dataReceived(self, data):
        self.received += len(data)

    def connectionLost(self, reason):
        assert self.received == 1_000_000, self.received

react(lambda reactor: loopbackTCP(Server(), Client(), noisy=True))