aewens
/
abots
1
0
Fork 1
Code Issues Pull Requests Releases Wiki Activity

Fixed socket to be more versitile

This commit is contained in:
aewens 2019-05-10 06:57:06 +02:00
parent c462ab1be3
commit f4fc9d2d01
7 changed files with 139 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
abots/helpers/__init__.py
View File

@ -1,7 +1,7 @@
from abots.helpers.hash import create_hash, md5, sha1, sha256, sha512
from abots.helpers.encode import jots, jsto, b64e, b64d, h2b64, b642h, ctos
from abots.helpers.numbers import clamp, randfloat
from abots.helpers.numbers import clamp, randfloat, isnumeric
from abots.helpers.general import eprint, deduce, noop, cast, get_digit, obtain
from abots.helpers.general import utc_now, utc_now_timestamp
from abots.helpers.logging import Logger
from abots.helpers.black_magic import infinitedict
from abots.helpers.black_magic import infinitedict

2
abots/helpers/general.py
View File

@ -32,4 +32,4 @@ def utc_now():
return datetime.utcnow().replace(tzinfo=timezone.utc)
def utc_now_timestamp():
return utc_now().timestamp()
return utc_now().timestamp()

5
abots/helpers/numbers.py
View File

@ -7,4 +7,7 @@ def randfloat(start, end=None):
if end is None:
end = start
start = 0
return (end - start) * random() + start
return (end - start) * random() + start
def isnumeric(test):
return test.replace(".", "", 1).isdigit()

4
abots/net/__init__.py
View File

@ -1,2 +1,4 @@
from abots.net.socket_server import SocketServer
from abots.net.socket_client import SocketClient
from abots.net.socket_client import SocketClient
from abots.net.irc_socket_client import IRCSocketClient
from abots.net.prefix_socket_client import PrefixSocketClient

50
abots/net/irc_socket_client.py Executable file
View File

@ -0,0 +1,50 @@
"""
IRC Socket Client
=================
Formatted to read and write IRC socket data
"""
from abots.net import SocketClient
from socket import timeout as sock_timeout
class IRCSocketClient(SocketClient):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def _format_message(self, message, *args):
if len(args) > 0:
formatted = message.format(*args)
else:
formatted = message
packaged = f"{formatted}\r\n"
return packaged.encode()
def _get_message(self, decode=True):
try:
packet = self.sock.recv(self.buffer_size)
except (BrokenPipeError, OSError) as e:
return None
return packet.decode() if decode else packet
# See rfc2812
#max_message_length = 512
#while True:#len(data) <= max_message_length:
# # Automatically break loop to prevent infinite loop
# # Allow at least twice the needed iterations to occur exiting loop
# # Force bufsize to cap out at buffer_size
# try:
# packet = self.sock.recv()#self.buffer_size)
# # The socket can either be broken or no longer open at all
# except (BrokenPipeError, OSError) as e:
# #if not isinstance(e, sock_timeout):
# # self._attempt_reconnect()
# continue
# #print(packet.decode())
# if len(packet) == 0:
# break
# data = data + packet
##print(data.decode())
#return data.decode() if decode else data

67
abots/net/prefix_socket_client.py Executable file
View File

@ -0,0 +1,67 @@
"""
Prefix Socket Client
====================
Formatted to read and write prefix headers to socket data
"""
from abots.net import SocketClient
from struct import pack, unpack
from socket import timeout as sock_timeout
class PrefixSocketClient(SocketClient):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def _format_message(self, message, *args):
if len(args) > 0:
formatted = message.format(*args)
else:
formatted = message
packaged = pack(">I", len(formatted)) + formatted.encode()
return packaged
def _recv_bytes(self, get_bytes, decode=True):
data = "".encode()
attempts = 0
while len(data) < get_bytes:
# Automatically break loop to prevent infinite loop
# Allow at least twice the needed iterations to occur exiting loop
if attempts > 2 * (get_bytes / self.buffer_size):
break
else:
attempts = attempts + 1
bufsize = get_bytes - len(data)
# Force bufsize to cap out at buffer_size
if bufsize > self.buffer_size:
bufsize = self.buffer_size
try:
packet = self.sock.recv(bufsize)
# The socket can either be broken or no longer open at all
except (BrokenPipeError, OSError) as e:
if not isinstance(e, sock_timeout):
self._attempt_reconnect()
return None
data = data + packet
return data.decode() if decode else data
def _get_message_size(self):
raw_message_size = self._recv_bytes(4, False)
if not raw_message_size:
return None
message_size = unpack(">I", raw_message_size)[0]
return message_size
def _get_message(self):
message_size = self._get_message_size()
if message_size is None:
return None
try:
return self._recv_bytes(message_size)
except OSError:
return None

50
abots/net/socket_client.py
View File

@ -83,54 +83,27 @@ class SocketClient(Thread):
self.reconnecting.wait()
self.send_message(message)
def _recv_bytes(self, get_bytes, decode=True):
def _get_message(self, decode=True):
data = "".encode()
attempts = 0
while len(data) < get_bytes:
# Automatically break loop to prevent infinite loop
# Allow at least twice the needed iterations to occur exiting loop
if attempts > 2 * (get_bytes / self.buffer_size):
break
else:
attempts = attempts + 1
bufsize = get_bytes - len(data)
# Force bufsize to cap out at buffer_size
if bufsize > self.buffer_size:
bufsize = self.buffer_size
while True:
try:
packet = self.sock.recv(bufsize)
packet = self.sock.recv(self.buffer_size)
# The socket can either be broken or no longer open at all
except (BrokenPipeError, OSError) as e:
if not isinstance(e, sock_timeout):
self._attempt_reconnect()
return None
if len(packet) == 0:
break
data = data + packet
return data.decode() if decode else data
def _package_message(self, message, *args):
def _format_message(self, message, *args):
if len(args) > 0:
formatted = message.format(*args)
else:
formatted = message
packaged = pack(">I", len(formatted)) + formatted.encode()
return packaged
def _get_message_size(self):
raw_message_size = self._recv_bytes(4, False)
if not raw_message_size:
return None
message_size = unpack(">I", raw_message_size)[0]
return message_size
def _get_message(self):
message_size = self._get_message_size()
if message_size is None:
return None
try:
return self._recv_bytes(message_size)
except OSError:
return None
return formatted.encode()
def _attempt_reconnect(self):
if self.kill_switch.is_set():
@ -165,16 +138,16 @@ class SocketClient(Thread):
self.stop()
def send_message(self, message, *args):
packaged = self._package_message(message, *args)
formatted = self._format_message(message, *args)
try:
self.sock.send(packaged)
self.sock.send(formatted)
except (BrokenPipeError, OSError) as e:
if not isinstance(e, sock_timeout):
self._attempt_reconnect()
self._attempt_reconnect()
def recv(self):
return [letter for letter in self._obtain(self._outbox)]
yield from self._obtain(self._outbox)
def send(self, message):
self._inbox.put(message)
@ -192,6 +165,7 @@ class SocketClient(Thread):
if self.broken.is_set():
self.reconnecting.wait()
message = self._get_message()
#print(message)
if message is None:
continue
self._outbox.put(message)
@ -208,4 +182,4 @@ class SocketClient(Thread):
self.sock.close()
self.stopped.set()
cast(done, "set")
# print("Stopped client!")
# print("Stopped client!")