diff options
Diffstat (limited to 'server')
| -rw-r--r-- | server/piztor/gen_users.py | 10 | ||||
| -rw-r--r-- | server/piztor/mesg_receiver.py | 20 | ||||
| -rw-r--r-- | server/piztor/mesg_sender.py | 23 | ||||
| -rw-r--r-- | server/piztor/prober.py | 196 | ||||
| -rw-r--r-- | server/piztor/ptp.rst | 37 | ||||
| -rw-r--r-- | server/piztor/ptp_send.py | 258 | ||||
| -rw-r--r-- | server/piztor/server.py | 201 |
7 files changed, 532 insertions, 213 deletions
diff --git a/server/piztor/gen_users.py b/server/piztor/gen_users.py new file mode 100644 index 0000000..2d4d67b --- /dev/null +++ b/server/piztor/gen_users.py @@ -0,0 +1,10 @@ +from random import randint + +def get_rand_gid(): + return randint(1, 2) + +def get_rand_sex(): + return randint(0, 1) + +for i in xrange(100): + print i, i, get_rand_gid(), get_rand_sex() diff --git a/server/piztor/mesg_receiver.py b/server/piztor/mesg_receiver.py new file mode 100644 index 0000000..d6fbf35 --- /dev/null +++ b/server/piztor/mesg_receiver.py @@ -0,0 +1,20 @@ +from random import random +from time import sleep + +from sys import argv +from ptp_send import * + +username = "a" +password = "a" +#username = "1234567890123456789012" +#password = "world12345678901234567890" + +if len(argv) == 2: + host = argv[1] + +if len(argv) == 3: + username = argv[1] + password = argv[2] + +uid, token = user_auth(username, password) +open_push_tunnel(token, username) diff --git a/server/piztor/mesg_sender.py b/server/piztor/mesg_sender.py new file mode 100644 index 0000000..e3ad6e4 --- /dev/null +++ b/server/piztor/mesg_sender.py @@ -0,0 +1,23 @@ +from random import random +from time import sleep + +from sys import argv +from ptp_send import * + +username = "a" +password = "a" +mesg = "niu x push!" +#username = "1234567890123456789012" +#password = "world12345678901234567890" + +if len(argv) == 2: + host = argv[1] + +if len(argv) == 4: + username = argv[1] + password = argv[2] + mesg = argv[3] + + +uid, token = user_auth(username, password) +send_text_mesg(token, username, mesg) diff --git a/server/piztor/prober.py b/server/piztor/prober.py index 9cfde24..00208cd 100644 --- a/server/piztor/prober.py +++ b/server/piztor/prober.py @@ -1,121 +1,8 @@ -import socket, logging -from struct import * from random import random -from select import select from time import sleep -FORMAT = "%(asctime)-15s %(message)s" -logging.basicConfig(format = FORMAT) -logger = logging.getLogger('piztor_server') -logger.setLevel(logging.INFO) - -def get_hex(data): - return "".join([hex(ord(c))[2:].zfill(2) for c in data]) - -class _SectionSize: - LENGTH = 4 - OPT_ID = 1 - STATUS = 1 - USER_ID = 4 - USER_TOKEN = 32 - GROUP_ID = 4 - ENTRY_CNT = 4 - LATITUDE = 8 - LONGITUDE = 8 - LOCATION_ENTRY = USER_ID + LATITUDE + LONGITUDE - PADDING = 1 - -host = "69.85.86.42" #"localhost" -port = 2223 - -def gen_auth(username, password): - length = _SectionSize.LENGTH + \ - _SectionSize.OPT_ID + \ - len(username) + \ - _SectionSize.PADDING + \ - len(password) + \ - _SectionSize.PADDING - - data = pack("!LB", length, 0x00) - data += username - data += "\0" - data += password - data += "\0" - return data - -def gen_auth_head_length(token, username): - return _SectionSize.USER_TOKEN + \ - len(username) + \ - _SectionSize.PADDING - - -def gen_update_location(token, username, lat, lng): - length = _SectionSize.LENGTH + \ - _SectionSize.OPT_ID + \ - gen_auth_head_length(token, username) + \ - _SectionSize.LATITUDE + \ - _SectionSize.LONGITUDE - - data = pack("!LB32s", length, 0x01, token) - data += username - data += chr(0) - data += pack("!dd", lat, lng) - return data - -def gen_request_location(token, username, gid): - length = _SectionSize.LENGTH + \ - _SectionSize.OPT_ID + \ - gen_auth_head_length(token, username) + \ - _SectionSize.GROUP_ID - - data = pack("!LB32s", length, 0x02, token) - data += username - data += chr(0) - data += pack("!L", gid) - return data - - -def gen_request_user_info(token, username, uid): - length = _SectionSize.LENGTH + \ - _SectionSize.OPT_ID + \ - gen_auth_head_length(token, username) + \ - _SectionSize.USER_ID - - data = pack("!LB32s", length, 0x03, token) - data += username - data += chr(0) - data += pack("!L", uid) - return data - -def gen_logout(token, username): - length = _SectionSize.LENGTH + \ - _SectionSize.OPT_ID + \ - gen_auth_head_length(token, username) - data = pack("!LB32s", length, 0x04, token) - data += username - data += chr(0) - return data - -def send(data): - received = bytes() - try: - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect((host, port)) - sock.sendall(data) - while True: - rd, wr, err = select([sock], [], [], 10) - if rd: - buff = sock.recv(4096) - if len(buff) == 0: - break - received += buff - else: - break - finally: - sock.close() - return received - from sys import argv +from ptp_send import * username = "a" password = "a" @@ -129,83 +16,6 @@ if len(argv) == 3: username = argv[1] password = argv[2] -def request_location(token, username, gid): - resp = send(gen_request_location(token, username, gid)) - try: - pl, optcode, status = unpack("!LBB", resp[:6]) - if pl != len(resp): - logger.error("Request location: incorrect packet length") - idx = 6 - while idx < pl: - uid, lat, lng = unpack("!Ldd", resp[idx:idx + 20]) - idx += 20 - print (uid, lat, lng) - except error: - logger.error("Request location: can not parse the response") - -def user_auth(username, password): - resp = send(gen_auth(username, password)) - try: - pl, optcode, status, uid, token = unpack("!LBBL32s", resp) - if pl != len(resp): - logger.error("User authentication: incorrect packet length") - print "status: " + str(status) - print "uid: " + str(uid) - print "token: " + get_hex(token) - except error: - logger.error("User authentication: can not parse the response") - - return uid, token - -def update_location(token, username, lat, lng): - resp = send(gen_update_location(token, username, lat, lng)) - print get_hex(resp) - try: - pl, optcode, status = unpack("!LBB", resp[:6]) - if pl != len(resp): - logger.error("Request location: incorrect packet length") - print "status: " + str(status) - except error: - logger.error("Request location: can not parse the response") - - -def request_user_info(token, username, uid): - resp = send(gen_request_user_info(token, username, uid)) - try: - pl, optcode, status = unpack("!LBB", resp[:6]) - if pl != len(resp): - logger.error("Request user info: incorrect packet length") - - idx = 6 - comp_id = None - sec_id = None - sex = None - while idx < pl: - info_key, = unpack("!B", resp[idx:idx + 1]) - idx += 1 - if info_key == 0x00: - gid, = unpack("!L", resp[idx:idx + 4]) - a, b, comp_id, sec_id = unpack("!BBBB", resp[idx:idx + 4]) - idx += 4 - print "gid: {}".format(str(gid)) - elif info_key == 0x01: - sex, = unpack("!B", resp[idx:idx + 1]) - idx += 1 - print "sex: {}".format(str(sex)) - return comp_id, sec_id, sex - except error: - logger.error("Request user info: can not parse the response") - -def logout(token, username): - resp = send(gen_logout(token, username)) - try: - pl, optcode, status = unpack("!LBB", resp) - if pl != len(resp): - logger.error("Logout: incorrect packet length") - print "status: " + str(status) - except error: - logger.error("Logout: can not parse the response") - for i in xrange(10): @@ -214,8 +24,8 @@ for i in xrange(10): comp_id, sec_id, sex = request_user_info(token, username, uid) if comp_id: - request_location(token, username, comp_id * 256 + sec_id) - request_location(token, username, comp_id * 256 + 0xff) + request_location(token, username, comp_id, sec_id) + request_location(token, username, comp_id, 0xff) logout(token, username) diff --git a/server/piztor/ptp.rst b/server/piztor/ptp.rst index a5523a8..1dc4f66 100644 --- a/server/piztor/ptp.rst +++ b/server/piztor/ptp.rst @@ -168,7 +168,7 @@ Piztor Transmission Protocol v1.0 | 0x04 | STATUS | +------+--------+ -- Push Tunnel +- Open Push Tunnel - Request @@ -186,20 +186,39 @@ Piztor Transmission Protocol v1.0 | 0x05 | STATUS | +------+--------+ -- Push Text Message +- Send Text Message - Request :: - - +--1b--+----?b----+ - | 0x10 | MESSAGE | - +------+--string--+ + + +--1b--+----?b-----+----?b----+ + | 0x06 | AUTH_HEAD | MESSAGE | + +------+-----------+--string--+ - Response :: - +--1b--+ - | 0x10 | - +------+ + +--1b--+---1b---+ + | 0x06 | STATUS | + +------+--------+ + +- Push + + - Acknowledgement + + :: + + +--1b--+-------32b---------+ + | 0x00 | PUSH_FINGERPRINT | + +------+-------------------+ + + - Text Message + + :: + + +--1b--+-------32b---------+----?b----+ + | 0x00 | PUSH_FINGERPRINT | MESSAGE | + +------+-------------------+--string--+ + diff --git a/server/piztor/ptp_send.py b/server/piztor/ptp_send.py new file mode 100644 index 0000000..b820041 --- /dev/null +++ b/server/piztor/ptp_send.py @@ -0,0 +1,258 @@ +from struct import * +import socket, logging +from select import select +from time import sleep + +FORMAT = "%(asctime)-15s %(message)s" +logging.basicConfig(format = FORMAT) +logger = logging.getLogger('piztor_server') +logger.setLevel(logging.INFO) + +def get_hex(data): + return "".join([hex(ord(c))[2:].zfill(2) for c in data]) + +class _SectionSize: + LENGTH = 4 + OPT_ID = 1 + STATUS = 1 + USER_ID = 4 + USER_TOKEN = 32 + GROUP_ID = 2 + ENTRY_CNT = 4 + LATITUDE = 8 + LONGITUDE = 8 + LOCATION_ENTRY = USER_ID + LATITUDE + LONGITUDE + PADDING = 1 + +host = "localhost" #"localhost" +port = 2222 + +def gen_auth(username, password): + length = _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + len(username) + \ + _SectionSize.PADDING + \ + len(password) + \ + _SectionSize.PADDING + + data = pack("!LB", length, 0x00) + data += username + data += "\0" + data += password + data += "\0" + return data + +def gen_auth_head_length(token, username): + return _SectionSize.USER_TOKEN + \ + len(username) + \ + _SectionSize.PADDING + + +def gen_update_location(token, username, lat, lng): + length = _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + gen_auth_head_length(token, username) + \ + _SectionSize.LATITUDE + \ + _SectionSize.LONGITUDE + + data = pack("!LB32s", length, 0x01, token) + data += username + data += chr(0) + data += pack("!dd", lat, lng) + return data + +def gen_request_location(token, username, comp_id, sec_id): + length = _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + gen_auth_head_length(token, username) + \ + _SectionSize.GROUP_ID + + data = pack("!LB32s", length, 0x02, token) + data += username + data += chr(0) + data += pack("!BB", comp_id, sec_id) + return data + + +def gen_request_user_info(token, username, uid): + length = _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + gen_auth_head_length(token, username) + \ + _SectionSize.USER_ID + + data = pack("!LB32s", length, 0x03, token) + data += username + data += chr(0) + data += pack("!L", uid) + return data + +def gen_logout(token, username): + length = _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + gen_auth_head_length(token, username) + data = pack("!LB32s", length, 0x04, token) + data += username + data += chr(0) + return data + +def gen_open_push_tunnel(token, username): + length = _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + gen_auth_head_length(token, username) + data = pack("!LB32s", length, 0x05, token) + data += username + data += chr(0) + return data + +def gen_send_text_mesg(token, username, mesg): + length = _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + gen_auth_head_length(token, username) + \ + len(mesg) + \ + _SectionSize.PADDING + data = pack("!LB32s", length, 0x06, token) + data += username + data += chr(0) + data += mesg + data += chr(0) + return data + +def send(data): + received = bytes() + try: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((host, port)) + sock.sendall(data) + while True: + rd, wr, err = select([sock], [], [], 10) + if rd: + buff = sock.recv(4096) + if len(buff) == 0: + break + received += buff + else: + break + finally: + sock.close() + return received + +def request_location(token, username, comp_id, sec_id): + resp = send(gen_request_location(token, username, comp_id, sec_id)) + try: + pl, optcode, status = unpack("!LBB", resp[:6]) + if pl != len(resp): + logger.error("Request location: incorrect packet length") + idx = 6 + while idx < pl: + uid, lat, lng = unpack("!Ldd", resp[idx:idx + 20]) + idx += 20 + print (uid, lat, lng) + except error: + logger.error("Request location: can not parse the response") + +def user_auth(username, password): + resp = send(gen_auth(username, password)) + try: + pl, optcode, status, uid, token = unpack("!LBBL32s", resp) + if pl != len(resp): + logger.error("User authentication: incorrect packet length") + print "status: " + str(status) + print "uid: " + str(uid) + print "token: " + get_hex(token) + except error: + logger.error("User authentication: can not parse the response") + + return uid, token + +def update_location(token, username, lat, lng): + resp = send(gen_update_location(token, username, lat, lng)) + print get_hex(resp) + try: + pl, optcode, status = unpack("!LBB", resp[:6]) + if pl != len(resp): + logger.error("Request location: incorrect packet length") + print "status: " + str(status) + except error: + logger.error("Request location: can not parse the response") + + +def request_user_info(token, username, uid): + resp = send(gen_request_user_info(token, username, uid)) + try: + pl, optcode, status = unpack("!LBB", resp[:6]) + if pl != len(resp): + logger.error("Request user info: incorrect packet length") + + idx = 6 + comp_id = None + sec_id = None + sex = None + while idx < pl: + info_key, = unpack("!B", resp[idx:idx + 1]) + idx += 1 + if info_key == 0x00: + comp_id, sec_id = unpack("!BB", resp[idx:idx + 2]) + idx += 2 + elif info_key == 0x01: + sex, = unpack("!B", resp[idx:idx + 1]) + idx += 1 + return comp_id, sec_id, sex + except error: + logger.error("Request user info: can not parse the response") + +def logout(token, username): + resp = send(gen_logout(token, username)) + try: + pl, optcode, status = unpack("!LBB", resp) + if pl != len(resp): + logger.error("Logout: incorrect packet length") + print "status: " + str(status) + except error: + logger.error("Logout: can not parse the response") + +def send_text_mesg(token, username, mesg): + resp = send(gen_send_text_mesg(token, username, mesg)) + try: + pl, optcode, status = unpack("!LBB", resp) + if pl != len(resp): + logger.error("Send text mesg: incorrect packet length") + print "status: " + str(status) + except error: + logger.error("Send text mesg: can not parse the response") + +def open_push_tunnel(token, username): + + length = _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + gen_auth_head_length(token, username) + data = pack("!LB32s", length, 0x05, token) + data += username + data += chr(0) + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((host, port)) + sock.sendall(data) + sock.recv(4096) + + length = -1 + while True: + received = bytes() + while True: + if len(received) > 4: + length, optcode = unpack("!LB", received[:5]) + + if len(received) == length: + break + rd, wr, err = select([sock], [], []) + if rd: + buff = sock.recv(4096) + if len(buff) == 0: + break + received += buff + else: + break + print len(received) + pl, optcode, fingerprint = unpack("!LB32s", received[:37]) + mesg = received[37:-1] + logger.info("Received a push: %s", mesg) + sock.sendall(pack("!LB32s", 37, 0x00, fingerprint)) diff --git a/server/piztor/server.py b/server/piztor/server.py index 3f1f2cb..bb88b53 100644 --- a/server/piztor/server.py +++ b/server/piztor/server.py @@ -7,6 +7,8 @@ from sqlalchemy import create_engine, and_ from sqlalchemy.orm import sessionmaker from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound +from collections import deque + import struct import os import logging @@ -45,18 +47,76 @@ _MAX_AUTH_HEAD_SIZE = _SectionSize.USER_TOKEN + \ _HEADER_SIZE = _SectionSize.LENGTH + \ _SectionSize.OPT_ID +_MAX_TEXT_MESG_SIZE = 1024 + class _OptCode: user_auth = 0x00 location_update = 0x01 location_info= 0x02 user_info = 0x03 user_logout = 0x04 + open_push_tunnel = 0x05 + send_text_mesg = 0x06 class _StatusCode: sucess = 0x00 failure = 0x01 +class PushData(object): + from hashlib import sha256 + def __init__(self, data): + self.data = data + self.finger_print = sha256(data).digest() + +class PushTextMesgData(PushData): + def __init__(self, mesg): + self.finger_print = sha256(mesg).digest() + buff = struct.pack("!B32s", 0x00, self.finger_print) + buff += mesg + buff += chr(0) + buff = struct.pack("!L", _SectionSize.LENGTH + len(buff)) + buff + self.data = buff + + +class PushTunnel(object): + def __init__(self): + self.pending = deque() + self.conn = None + + def __del__(self): + if self.conn: + self.conn.loseConnection() + + def add(self, pdata): + logger.info("-- Push data enqued --") + logger.info("Data: %s", get_hex(pdata.data)) + self.pending.append(pdata) + + def on_receive(self, data): + front = self.pending.popleft() + length, optcode, fingerprint = struct.unpack("!LB32s", data) + if front.finger_print != fingerprint: + raise PiztorError + logger.info("-- Push data confirmed --") + self.push() + + def push(self): + if (self.conn is None) or len(self.pending) == 0: + return + print "Pushing" + front = self.pending.popleft() + self.pending.appendleft(front) + self.conn.transport.write(front.data) + + def connect(self, conn): + conn.tunnel = self + self.conn = conn + + def on_connection_lost(self): + self.conn = None + class RequestHandler(object): + push_tunnels = dict() def __init__(self): Session = sessionmaker(bind = engine) self.session = Session() @@ -78,7 +138,9 @@ class RequestHandler(object): if uauth.user.username != username: logger.warning("Toke and username mismatch") return None - + uid = uauth.uid + if not cls.push_tunnels.has_key(uid): + cls.push_tunnels[uid] = PushTunnel() return uauth except NoResultFound: @@ -122,7 +184,7 @@ class UserAuthHandler(RequestHandler): bytes('\x00' * 32)) - def handle(self, tr_data): + def handle(self, tr_data, conn): self.check_size(tr_data) logger.info("Reading auth data...") pos = -1 @@ -178,7 +240,7 @@ class LocationUpdateHandler(RequestHandler): _SectionSize.OPT_ID + \ _SectionSize.STATUS - def handle(self, tr_data): + def handle(self, tr_data, conn): self.check_size(tr_data) logger.info("Reading location update data...") try: @@ -224,7 +286,7 @@ class LocationInfoHandler(RequestHandler): _SectionSize.STATUS + \ _SectionSize.LOCATION_ENTRY * item_num - def handle(self, tr_data): + def handle(self, tr_data, conn): self.check_size(tr_data) logger.info("Reading location request data..") try: @@ -298,7 +360,7 @@ class UserInfoHandler(RequestHandler): info_key = entry_code return struct.pack("!B", info_key) + pack_method(user) - def handle(self, tr_data): + def handle(self, tr_data, conn): self.check_size(tr_data) logger.info("Reading user info request data...") try: @@ -348,7 +410,7 @@ class UserLogoutHandler(RequestHandler): _SectionSize.OPT_ID + \ _SectionSize.STATUS - def handle(self, tr_data): + def handle(self, tr_data, conn): self.check_size(tr_data) logger.info("Reading user logout data...") try: @@ -368,15 +430,113 @@ class UserLogoutHandler(RequestHandler): if uauth is None: logger.warning("Authentication failure") return struct.pack("!LBB", self._response_size, - _OptCode.location_update, + _OptCode.user_logout, _StatusCode.failure) + del RequestHandler.push_tunnels[uauth.uid] uauth.regen_token() logger.info("User Logged out successfully!") self.session.commit() return struct.pack("!LBB", self._response_size, _OptCode.user_logout, _StatusCode.sucess) - + +class OpenPushTunnelHandler(RequestHandler): + + _max_tr_data_size = _MAX_AUTH_HEAD_SIZE + + _response_size = \ + _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + _SectionSize.STATUS + + def handle(self, tr_data, conn): + self.check_size(tr_data) + logger.info("Reading open push tunnel data...") + try: + token, = struct.unpack("!32s", tr_data[:32]) + username, tail = RequestHandler.trunc_padding(tr_data[32:]) + if username is None: + raise struct.error + except struct.error: + raise BadReqError("Open push tunnel: Malformed request body") + + logger.info("Trying to open push tunnel with " + "(token = {0}, username = {1})"\ + .format(get_hex(token), username)) + + uauth = RequestHandler.get_uauth(token, username, self.session) + # Authentication failure + if uauth is None: + logger.warning("Authentication failure") + return struct.pack("!LBB", self._response_size, + _OptCode.open_push_tunnel, + _StatusCode.failure) + + tunnel = RequestHandler.push_tunnels[uauth.uid] + pt = RequestHandler.push_tunnels + uid = uauth.uid + if pt.has_key(uid): + tunnel = pt[uid] + tunnel.connect(conn) + tunnel.push() + + logger.info("Push tunnel opened successfully!") + return struct.pack("!LBB", self._response_size, + _OptCode.open_push_tunnel, + _StatusCode.sucess) + +class SendTextMessageHandler(RequestHandler): + + _max_tr_data_size = _MAX_AUTH_HEAD_SIZE + \ + _MAX_TEXT_MESG_SIZE + \ + _SectionSize.PADDING + + _response_size = \ + _SectionSize.LENGTH + \ + _SectionSize.OPT_ID + \ + _SectionSize.STATUS + + def handle(self, tr_data, conn): + self.check_size(tr_data) + logger.info("Reading send text mesg data...") + try: + token, = struct.unpack("!32s", tr_data[:32]) + username, tail = RequestHandler.trunc_padding(tr_data[32:]) + mesg = tail[:-1] + if username is None: + raise struct.error + except struct.error: + raise BadReqError("Send text mesg: Malformed request body") + + logger.info("Trying to send text mesg with " + "(token = {0}, username = {1})"\ + .format(get_hex(token), username)) + + uauth = RequestHandler.get_uauth(token, username, self.session) + # Authentication failure + if uauth is None: + logger.warning("Authentication failure") + return struct.pack("!LBB", self._response_size, + _OptCode.send_text_mesg, + _StatusCode.failure) + + pt = RequestHandler.push_tunnels + u = uauth.user + ulist = self.session.query(UserModel) \ + .filter(and_(UserModel.comp_id == u.comp_id, + UserModel.sec_id == u.sec_id)).all() + + for user in ulist: + uid = user.id + if pt.has_key(uid): + tunnel = pt[uid] + tunnel.add(PushTextMesgData(mesg)) + tunnel.push() + logger.info("Sent text mesg successfully!") + return struct.pack("!LBB", self._response_size, + _OptCode.send_text_mesg, + _StatusCode.sucess) + class PTP(Protocol, TimeoutMixin): @@ -384,7 +544,9 @@ class PTP(Protocol, TimeoutMixin): LocationUpdateHandler, LocationInfoHandler, UserInfoHandler, - UserLogoutHandler] + UserLogoutHandler, + OpenPushTunnelHandler, + SendTextMessageHandler] handler_num = len(handlers) @@ -399,9 +561,11 @@ class PTP(Protocol, TimeoutMixin): self.buff = bytes() self.length = -1 self.factory = factory + self.tunnel = None def timeoutConnection(self): logger.info("The connection times out") + self.transport.loseConnection() def connectionMade(self): logger.info("A new connection is made") @@ -425,21 +589,36 @@ class PTP(Protocol, TimeoutMixin): print self.length, PTP._MAX_REQUEST_SIZE raise BadReqError("The size of remaining part is too big") if len(self.buff) == self.length: + if self.tunnel: # received push response + self.tunnel.on_receive(self.buff) + self.buff = bytes() + self.length = -1 + return h = PTP.handlers[self.optcode]() - reply = h.handle(self.buff[5:]) + reply = h.handle(self.buff[5:], self) logger.info("Wrote: %s", get_hex(reply)) self.transport.write(reply) + if self.tunnel: + logger.info("Blocking the client...") + self.buff = bytes() + self.length = -1 + self.setTimeout(None) + return self.transport.loseConnection() elif len(self.buff) > self.length: raise BadReqError("The actual length is larger than promised") except BadReqError as e: logger.warn("Rejected a bad request: %s", str(e)) + self.transport.loseConnection() except DBCorruptionError: logger.error("*** Database corruption ***") - finally: + self.transport.loseConnection() + if self.tunnel is None: self.transport.loseConnection() def connectionLost(self, reason): + if self.tunnel: + self.tunnel.on_connection_lost() logger.info("The connection is lost") self.setTimeout(None) |