class RemoteConnection(asyncore.dispatcher):
buffer_send_raw = b''
buffer_send = b''
buffer_recv_raw = b''
buffer_recv = b''
local = None
cipher = None
def __init__(self, host, port, token):
asyncore.dispatcher.__init__(self)
self.buffer_send_raw += token
magic = bytearray(os.urandom(MAGIC_LEN))
self.buffer_send_raw += magic
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
self.connect((host, port))
self.cipher = AESCipher()
def handle_connect(self):
pass
def handle_close(self):
LOGGER.info('%s remote close', id(self))
self.local.handle_close()
self.close()
def handle_read(self):
data = self.recv(BUF_SIZE)
#LOGGER.debug('%s remote recv: %s', id(self), data)
self.buffer_recv += data
ddata, dlen = self.cipher.decrypt_all(self.buffer_recv)
self.buffer_recv_raw += ddata
self.buffer_recv = self.buffer_recv[dlen:]
if len(self.buffer_recv_raw) > 0:
self.local.buffer_send += self.buffer_recv_raw
self.buffer_recv_raw = b''
def writable(self):
if len(self.buffer_send_raw) > 0:
self.buffer_send += self.cipher.encrypt_all(self.buffer_send_raw)
self.buffer_send_raw = b''
return (len(self.buffer_send) > 0)
def handle_write(self):
sent = self.send(self.buffer_send)
#LOGGER.debug('%s remote send %s', id(self), self.buffer[:sent])
self.buffer_send = self.buffer_send[sent:]
class LocalConnection(asyncore.dispatcher):
buffer_recv = b''
buffer_recv_raw = b''
buffer_send = b''
buffer_send_raw = b''
server = None
remote = None
stage = STAGE_INIT
cipher = None
token = None
def __init__(self, sock):
super().__init__(sock)
self.cipher = AESCipher()
def handle_read(self):
data = self.recv(BUF_SIZE)
if not data:
return
if self.token:
self.server.add_traffic(self.token, len(data) / (1024.0 * 1024.0))
self.buffer_recv += data
ddata, dlen = self.cipher.decrypt_all(self.buffer_recv)
self.buffer_recv_raw += ddata
self.buffer_recv = self.buffer_recv[dlen:]
#LOGGER.debug('%s local recv %s', id(self), data)
while True:
if self.stage == STAGE_INIT:
if len(self.buffer_recv_raw) < TOKEN_LEN + MAGIC_LEN:
return
token = self.buffer_recv_raw[0:TOKEN_LEN]
self.token = token
magic = self.buffer_recv_raw[TOKEN_LEN:TOKEN_LEN + MAGIC_LEN]
LOGGER.info("%s accept token: %s, magic: %s", id(self),
token.hex(), magic.hex())
# auth
if not self.server.query_token(token):
LOGGER.warning('%s token not exist!', id(self))
self.handle_close()
return
if not self.server.query_magic(magic):
LOGGER.warning('%s duplicated magic!', id(self))
self.handle_close()
return
self.buffer_recv_raw = self.buffer_recv_raw[TOKEN_LEN +
MAGIC_LEN:]
self.stage = STAGE_HANDSHAKE
continue
elif self.stage == STAGE_HANDSHAKE:
if len(self.buffer_recv_raw) < 1:
return
atyp = self.buffer_recv_raw[0]
addr = ""
port = 0
if atyp == 1:
if len(self.buffer_recv_raw) < 7:
return
addr = str(self.buffer_recv_raw[1]) + '.' + str(self.buffer_recv_raw[2]) + \
'.' + str(self.buffer_recv_raw[3]) + '.' + str(self.buffer_recv_raw[4])
port = self.buffer_recv_raw[
5] * 256 + self.buffer_recv_raw[6]
self.buffer_recv_raw = self.buffer_recv_raw[7:]
elif atyp == 3:
if len(self.buffer_recv_raw) < 2:
return
alen = self.buffer_recv_raw[1]
if len(self.buffer_recv_raw) < 2 + alen + 2:
return
addr = self.buffer_recv_raw[2:2 + alen].decode('utf-8')
port = self.buffer_recv_raw[
2 + alen] * 256 + self.buffer_recv_raw[2 + alen + 1]
self.buffer_recv_raw = self.buffer_recv_raw[2 + alen + 2:]
LOGGER.info('%s local handshake: %s:%d', id(self), addr, port)
try:
self.remote = RemoteConnection(addr, port)
except Exception:
LOGGER.error('%s cannot connect to %s:%d', id(self), addr,
port)
return
self.remote.local = self
self.remote.buffer += self.buffer_recv_raw
self.buffer_recv_raw = b''
self.stage = STAGE_STREAM
continue
elif self.stage == STAGE_STREAM:
self.remote.buffer += self.buffer_recv_raw
self.buffer_recv_raw = b''
return
def writable(self):
if len(self.buffer_send_raw) > 0:
self.buffer_send += self.cipher.encrypt_all(self.buffer_send_raw)
self.buffer_send_raw = b''
return (len(self.buffer_send) > 0 and self.stage == STAGE_STREAM)
def handle_write(self):
sent = self.send(self.buffer_send)
#LOGGER.debug('%s local send %s', id(self), self.buffer[:sent])
self.buffer_send = self.buffer_send[sent:]
if self.token:
self.server.add_traffic(self.token, sent / (1024.0 * 1024.0))
def handle_close(self):
LOGGER.info('%s local close', id(self))
if self in self.server.conn_list:
self.server.conn_list.remove(self)
self.close()