from twisted.internet import reactor, protocol
from Crypto.Cipher import AES

PORT = 8005
KEY = "@Ecb_aTTack>esew"
SECRET = "flag{ese@qq!hooray}"
def encrypt_block(key, plaintext):
        encobj = AES.new(key, AES.MODE_ECB)
        return encobj.encrypt(plaintext).encode('hex')
def encrypt(key, ptxt):
    if (len(ptxt) % 16 != 0):
        ptxt = ptxt + "*" * (16 - (len(ptxt) % 16))
    return encrypt_block(key, ptxt)
class Server(protocol.Protocol):

    def connectionMade(self):
        msg = '''
        Welcome to the AES_ECB Cryptography System
        To encrypt, enter "encrypt [plaintext]". Ex, encrypt 123
        1. cliphertext = plaintext + SECRET
        2. this padding is *      
        \r\n>'''
        self.transport.write(msg)

    def dataReceived(self, data):
        try:
            data = data.strip().split(" ")
            print (data)
            if data[0] == "encrypt":
                response = encrypt(KEY,data[1] + SECRET)
            else:
                raise
            self.transport.write(">" + response + "\r\n>")
        except:
            self.transport.write(">Something is wrong here\r\n>")
def main():
    factory = protocol.ServerFactory()
    factory.protocol = Server
    reactor.listenTCP(PORT, factory)
    reactor.run()
if __name__ == '__main__':
    main()