SMTP proxy to implement two factor auth

smtp_proxy.py

#encoding=utf-8
__author__ = 'lijiejie'

'''
在smtp认证的时候，用户名全部变成username.iqiyi.com的形式，否则连接主动被服务器断开
'''


import socket
import select
import time
import sys
import base64

buffer_size = 4096
delay = 0.0001
forward_to = ('applesmtp.163.com', 25)

class Forward:
    def __init__(self):
        self.forward = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    def start(self, host, port):
        try:
            self.forward.connect((host, port))
            return self.forward
        except Exception, e:
            print e
            return False


class TheServer:
    input_list = []
    channel = {}

    def __init__(self, host, port):
        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server.bind((host, port))
        self.server.listen(200)
        self.get_user_input = False

    def main_loop(self):
        self.input_list.append(self.server)
        while 1:
            time.sleep(delay)
            ss = select.select
            inputready, outputready, exceptready = ss(self.input_list, [], [])
            for self.s in inputready:
                if self.s == self.server:
                    self.on_accept()
                    break

                self.data = self.s.recv(buffer_size)
                if len(self.data) == 0:
                    self.on_close()
                    break
                else:
                    self.on_recv()

    def on_accept(self):
        forward = Forward().start(forward_to[0], forward_to[1])
        clientsock, clientaddr = self.server.accept()
        if forward:
            print clientaddr, "has connected"
            self.input_list.append(clientsock)
            self.input_list.append(forward)
            self.channel[clientsock] = forward
            self.channel[forward] = clientsock
        else:
            print "Can't establish connection with remote server.",
            print "Closing connection with client side", clientaddr
            clientsock.close()

    def on_close(self):
        print self.s.getpeername(), "has disconnected"
        self.input_list.remove(self.s)
        self.input_list.remove(self.channel[self.s])
        out = self.channel[self.s]
        self.channel[out].close()
        self.channel[self.s].close()
        del self.channel[out]
        del self.channel[self.s]

    def on_recv(self):
        data = self.data
        print data

        if self.get_user_input:
            self.username += data
        else:
            self.channel[self.s].send(data)

        if ord(data[-1]) == 10 and self.get_user_input:
            self.get_user_input = False
            print 'User input name:', base64.b64decode(self.username)
            if not base64.b64decode(self.username).lower().find('.iqiyi.com') > 0:
                self.channel[forward].close()
            self.username = base64.b64encode(base64.b64decode(self.username).strip().strip('.iqiyi.com'))
            print 'Send name', base64.b64decode(self.username)
            self.channel[self.s].send(self.username)
            self.channel[self.s].send(data)

        if data.strip() == '334 dXNlcm5hbWU6':    # username start
            self.username = ''
            self.get_user_input = True

if __name__ == '__main__':
        server = TheServer('', 25)
        try:
            server.main_loop()
        except KeyboardInterrupt:
            print "Ctrl C - Stopping server"
            sys.exit(1)