Reformat connection classes to accommodate proxy settings

This commit is contained in:
Dan 2018-01-16 16:28:50 +01:00
parent 0aed7bf24a
commit e9f6bce579
5 changed files with 15 additions and 21 deletions

View File

@ -56,7 +56,7 @@ class Connection:
def send(self, data: bytes):
with self.lock:
self.connection.send(data)
self.connection.sendall(data)
def recv(self) -> bytes or None:
return self.connection.recv()
return self.connection.recvall()

View File

@ -28,12 +28,6 @@ class TCP(socks.socksocket):
def __init__(self):
super().__init__()
def send(self, *args):
pass
def recv(self, *args):
pass
def close(self):
try:
self.shutdown(socket.SHUT_RDWR)

View File

@ -33,7 +33,7 @@ class TCPAbridged(TCP):
self.is_first_packet = True
log.info("Connected!")
def send(self, data: bytes):
def sendall(self, data: bytes, *args):
length = len(data) // 4
data = (
@ -48,20 +48,20 @@ class TCPAbridged(TCP):
super().sendall(data)
def recv(self) -> bytes or None:
length = self.recvall(1)
def recvall(self, length: int = 0) -> bytes or None:
length = super().recvall(1)
if length is None:
return None
if length == b"\x7f":
length = self.recvall(3)
length = super().recvall(3)
if length is None:
return None
length = int.from_bytes(length, "little") * 4
packet = self.recvall(length)
packet = super().recvall(length)
return packet

View File

@ -35,7 +35,7 @@ class TCPFull(TCP):
self.seq_no = 0
log.info("Connected!")
def send(self, data: bytes):
def sendall(self, data: bytes, *args):
# 12 = packet_length (4), seq_no (4), crc32 (4) (at the end)
data = pack("<II", len(data) + 12, self.seq_no) + data
data += pack("<I", crc32(data))
@ -43,13 +43,13 @@ class TCPFull(TCP):
super().sendall(data)
def recv(self) -> bytes or None:
length = self.recvall(4)
def recvall(self, length: int = 0) -> bytes or None:
length = super().recvall(4)
if length is None:
return None
packet = self.recvall(unpack("<I", length)[0] - 4)
packet = super().recvall(unpack("<I", length)[0] - 4)
if packet is None:
return None

View File

@ -34,7 +34,7 @@ class TCPIntermediate(TCP):
self.is_first_packet = True
log.info("Connected!")
def send(self, data: bytes):
def sendall(self, data: bytes, *args):
length = len(data)
data = pack("<i", length) + data
@ -44,12 +44,12 @@ class TCPIntermediate(TCP):
super().sendall(data)
def recv(self) -> bytes or None:
length = self.recvall(4)
def recvall(self, length: int = 0) -> bytes or None:
length = super().recvall(4)
if length is None:
return None
packet = self.recvall(unpack("<I", length)[0])
packet = super().recvall(unpack("<I", length)[0])
return packet