mirror of
https://github.com/TeamPGM/pyrogram.git
synced 2024-11-16 12:51:18 +00:00
Reformat connection classes to accommodate proxy settings
This commit is contained in:
parent
0aed7bf24a
commit
e9f6bce579
@ -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()
|
||||
|
@ -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)
|
||||
|
@ -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
|
||||
|
@ -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
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user