class TorSocketLoop(threading.Thread):
def __init__(self, our_sock, send_func):
super().__init__(name='SocketLoop{:x}'.format(our_sock.fileno()))
self._our_sock = our_sock
self._send_func = send_func
self._do_loop = True
self._cntrl_r, self._cntrl_w = socket.socketpair()
self._selector = DefaultSelector()
self._selector.register(self._our_sock, EVENT_READ, self._do_recv)
self._selector.register(self._cntrl_r, EVENT_READ, self._do_stop)
def _do_recv(self, sock):
try:
data = sock.recv(1024)
self._send_func(data)
except ConnectionResetError:
logger.debug('Client was badly disconnected...')
def _do_stop(self, sock):
self._do_loop = False
def _cleanup(self):
self._selector.unregister(self._cntrl_r)
self._cntrl_w.close()
self._cntrl_r.close()
self.close_sock()
self._selector.close()
@property
def fileno(self):
if not self._our_sock:
return None
return self._our_sock.fileno()
def close_sock(self):
if not self._our_sock:
return
self._selector.unregister(self._our_sock)
# self._our_sock.shutdown(socket.SHUT_WR)
self._our_sock.close()
self._our_sock = None
def stop(self):
self._cntrl_w.send(b'\1')
def run(self):
logger.debug('Starting...')
while self._do_loop:
events = self._selector.select()
for key, _ in events:
callback = key.data
callback(key.fileobj)
self._cleanup()
logger.debug('Stopped...')
def append(self, data):
self._our_sock.send(data)
class EventLoop:
def __init__(self, server_id, handler: Handler):
self.select = DefaultSelector()
self.server_id = server_id
self.handler = handler
def register(self, fd: Socket, events: SelectorKey, data=None):
fd.setblocking(False)
self.select.register(fd, events, data)
def modify(self, fd: Socket, events: SelectorKey, data=None):
self.select.modify(fd, events, data)
def unregister(self, fd):
self.select.unregister(fd)
def run(self):
self.register(self.server_id, selectors.EVENT_READ)
while True:
events = self.select.select()
for key, mask in events:
if mask == EVENT_READ:
if key.fileobj == self.server_id:
self.accept(key.fileobj)
else:
self.read(key.fileobj)
elif mask == EVENT_WRITE:
print(f"EventLoop, key={key}, data={key.data}")
self.write(key.fileobj, key.data)
def accept(self, sock):
"""sock 要处理的 fd, mask 为事件类型 READ or WRITE"""
conn = self.handler.accept(sock) # Should be ready
self.register(conn, EVENT_READ, self.read)
def read(self, conn):
for data in self.handler.read(conn):
if data and data not in ('q', 'Q'):
resp = self.handler.handle(conn, data)
self.modify(conn, EVENT_WRITE, resp)
else:
self.unregister(conn)
self.handler.close(conn)
def write(self, conn, data):
self.handler.send(conn, data)
self.modify(conn, EVENT_READ)
def close(self):
self.select.close()
class TorSocketLoop(threading.Thread):
def __init__(self, our_sock, send_func):
super().__init__(name='SocketLoop{:x}'.format(our_sock.fileno()))
self._our_sock = our_sock
self._send_func = send_func
self._do_loop = True
self._cntrl_r, self._cntrl_w = socket.socketpair()
self._selector = DefaultSelector()
self._selector.register(self._our_sock, EVENT_READ, self._do_recv)
self._selector.register(self._cntrl_r, EVENT_READ, self._do_stop)
def _do_recv(self, sock):
data = sock.recv(1024)
self._send_func(data)
def _do_stop(self, sock):
self._do_loop = False
def _cleanup(self):
self._selector.unregister(self._cntrl_r)
self._cntrl_w.close()
self._cntrl_r.close()
self._selector.unregister(self._our_sock)
self._our_sock.shutdown(socket.SHUT_WR)
self._our_sock.close()
self._selector.close()
def stop(self):
self._cntrl_w.send(b'\1')
def run(self):
logger.debug("Starting...")
while self._do_loop:
events = self._selector.select()
for key, _ in events:
callback = key.data
callback(key.fileobj)
self._cleanup()
logger.debug("Stopped...")
def append(self, data):
self._our_sock.send(data)
class Selector:
def __init__(self):
self._sel = DefaultSelector()
self._fds = {}
def register_read(self, task, fileno):
return self._register(task, fileno, EVENT_READ)
def register_write(self, task, fileno):
return self._register(task, fileno, EVENT_WRITE)
def _unregister(self, fd):
fileno = fd.fileno()
if fileno in self._fds:
LOG.debug('unregister fd %r', fd)
del self._fds[fileno]
self._sel.unregister(fileno)
def _register(self, task, fileno, events):
fd = self._fds.get(fileno, None)
if fd is None:
fd = KernelFd(self, fileno, task)
self._fds[fileno] = fd
if fd.task is not task:
raise RuntimeError(
'file descriptor already registered by other task')
if fd.selector_key is None:
LOG.debug('register fd %r, events=%r', fd, events)
fd.selector_key = self._sel.register(fd, events)
elif fd.events != events:
LOG.debug('modify fd %r, events=%r', fd, events)
fd.selector_key = self._sel.modify(fd, events)
return fd
def poll(self, timeout):
io_events = self._sel.select(timeout=timeout)
for key, mask in io_events:
fd = key.fileobj
if fd.task.is_alive and fd.events & mask:
LOG.debug('task %r wakeup by fd %r', fd.task, fd)
yield fd.task
def close(self):
self._sel.close()
def httpmcsleep():
sock4 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock6 = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
sock4.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
sock6.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
sock4.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, True)
sock6.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, True)
sock4.bind(("0.0.0.0", PORT))
sock6.bind(("::", PORT))
sock4.listen(128)
sock6.listen(128)
sock4.setblocking(False)
sock6.setblocking(False)
selector = DefaultSelector()
selector.register(sock4, EVENT_READ, handler_accept)
selector.register(sock6, EVENT_READ, handler_accept)
try:
while True:
for key, event in selector.select():
conn = key.fileobj
func = key.data
func(conn, selector)
except Exception as e:
print(f"httpmcsleep() 异常： {e}")
raise e
finally:
selector.close()
sock4.close()
sock6.close()
class HotKey:
"""
监听键盘、鼠标事件，触发动作。
"""
def __init__(self, device=None):
"""
device: /dev/input/eventX, default: all keyboard.
"""
self._hotkey_list = []
self._hotkey_seq_dict = {}
self.devices = []
mouses, kbms = libkbm.getkbm()
if len(kbms) <= 0:
raise HotKeyError("没有发现至少一个键盘。")
if device is None:
self.kbms = kbms
else:
if device not in kbms:
raise ValueError(f"{device} 不是键盘设备。")
else:
self.kbms = device
self.LISTEN = 1
self.REPLACE = 2
self._mode = self.LISTEN
# 把device注册到selectors
self._selector = DefaultSelector()
self._fileobjs = []
for device in self.kbms:
fd = open(device, "rb")
devfd = Device(fd)
self._selector.register(devfd.fd, EVENT_READ, devfd)
self._fileobjs.append(fd)
@property
def mode(self):
return self._mode
@mode.setter
def mode(self, m):
"""
m: replace mode: self.REPLACE, listen mode: self.LISTEN (default)
"""
if m == self.REPLACE or m == self.LISTEN:
self._mode = m
else:
raise ValueError("mode is choice: self.REPLACE or self.LISTEN")
def addhotkey(self, keyseq=(), callback=print):
"""
keyseq: ["alt", "f"]
callback: function()
"""
if not isinstance(keyseq, tuple):
raise HotKeyError("键序列必须是 tuple !")
# 预处理
seq = []
for key in keyseq:
key = key.upper()
if key == "ALT":
key = "LEFTALT"
elif key == "CTRL":
key = "LEFTCTRL"
key = ev.evbit("KEY_" + key)
seq.append(key)
seq_len = len(seq)
# 判断是否是最后一个键
last_key = 0
current_keyseq = self._hotkey_seq_dict
for key in seq:
last_key += 1
if hasattr(current_keyseq, "__call__"):
raise HotKeyError(f"{seq} 的父键已存在！")
if key in current_keyseq:
if last_key < seq_len:
current_keyseq = current_keyseq[key]
elif last_key == seq_len and key in current_keyseq:
raise HotKeyError(f"{seq} 按键序列已存在！")
else:
current_keyseq[key] = callback
self._hotkey_list.append(seq)
else:
if last_key == seq_len:
current_keyseq[key] = callback
self._hotkey_list.append(seq)
else:
# 一个新的快捷键序列
current_keyseq[key] = {}
current_keyseq = current_keyseq[key]
def watch(self):
"""
return: callback function
"""
while True:
for key, event_ in self._selector.select():
logger.debug("self._selector.select()")
devfd = key.data
hotkey = self._hotkey_seq_dict
try:
for e in devfd.events():
if e.matches(ev.evbit("EV_KEY")):
logger.debug(
f"key: {e.code.name} value: {e.value}")
logger.debug(f"hotkey: {hotkey}")
logger.debug("-" * 60)
# 这个事件在hotkey seq
if e.value == 1 and e.code in hotkey:
hotkey = hotkey.get(e.code)
if hasattr(hotkey, "__call__"):
return hotkey
elif e.value == 0 and e.code in hotkey:
hotkey = self._hotkey_seq_dict
except EventsDroppedException:
logger.warning("EventsDroppedException")
for e in devfd.sync():
logger.debug(e)
logger.debug("从这里返回的？？？？")
def watchrun(self):
func = self.watch()
func()
def close(self):
for fileobj in self._fileobjs:
if not fileobj.closed:
fileobj.close()
self._selector.close()
def __del__(self):
self.close()
class DBusConnection:
def __init__(self, sock):
self.sock = sock
self.parser = Parser()
self.outgoing_serial = count(start=1)
self.selector = DefaultSelector()
self.select_key = self.selector.register(sock, EVENT_READ)
self._stop_r, self._stop_w = os.pipe()
self.stop_key = self.selector.register(self._stop_r, EVENT_READ)
self.send_lock = Lock()
self.rcv_lock = Lock()
self.unique_name = None
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
return False
def send(self, message: Message, serial=None):
"""Serialise and send a :class:`~.Message` object"""
if serial is None:
serial = next(self.outgoing_serial)
data = message.serialise(serial=serial)
with self.send_lock:
self.sock.sendall(data)
def receive(self, *, timeout=None) -> Message:
"""Return the next available message from the connection
If the data is ready, this will return immediately, even if timeout<=0.
Otherwise, it will wait for up to timeout seconds, or indefinitely if
timeout is None. If no message comes in time, it raises TimeoutError.
If the connection is closed from another thread, this will raise
ReceiveStopped.
"""
if timeout is not None:
deadline = time.monotonic() + timeout
else:
deadline = None
with self.rcv_lock:
while True:
msg = self.parser.get_next_message()
if msg is not None:
return msg
if deadline is not None:
timeout = deadline - time.monotonic()
b = self._read_some_data(timeout)
self.parser.add_data(b)
def _read_some_data(self, timeout=None):
# Wait for data or a signal on the stop pipe
for key, ev in self.selector.select(timeout):
if key == self.select_key:
return unwrap_read(self.sock.recv(4096))
elif key == self.stop_key:
raise ReceiveStopped(
"DBus receive stopped from another thread")
raise TimeoutError
def interrupt(self):
"""Make any threads waiting for a message raise ReceiveStopped"""
os.write(self._stop_w, b'a')
def reset_interrupt(self):
"""Allow calls to .receive() again after .interrupt()
To avoid race conditions, you should typically wait for threads to
respond (e.g. by joining them) between interrupting and resetting.
"""
# Clear any data on the stop pipe
while (self.stop_key, EVENT_READ) in self.selector.select(timeout=0):
os.read(self._stop_r, 1024)
def close(self):
"""Close the connection"""
self.interrupt()
self.selector.close()
self.sock.close()
class TorReceiver(threading.Thread):
def __init__(self, tor_socket, handler_mgr):
super().__init__(name='RecvLoop_{}'.format(tor_socket.ip_address[0:7]))
self._tor_socket = tor_socket
self._handler_mgr = handler_mgr
self._do_loop = False
# fmt: off
self._regs_funcs_map = {
'reg': {
socket.socket: self.register_socket,
TorStream: self.register_stream
},
'unreg': {
socket.socket: self.unregister_socket,
TorStream: self.unregister_stream
}
}
# fmt: on
self._stream_to_callback = {}
self._selector = DefaultSelector()
self._cntrl_r, self._cntrl_w = socket.socketpair()
self._selector.register(self._cntrl_r, EVENT_READ, self._do_stop)
self._selector.register(self._tor_socket.ssl_socket, EVENT_READ, self._do_recv)
def _cleanup(self):
self._selector.unregister(self._cntrl_r)
self._cntrl_w.close()
self._cntrl_r.close()
self._selector.unregister(self._tor_socket.ssl_socket)
self._selector.close()
def start(self):
self._do_loop = True
super().start()
def stop(self):
logger.debug('Stopping receiver thread...')
self._cntrl_w.send(b'\1')
self.join()
def register(self, sock_or_stream, events, callback):
func = self._regs_funcs_map['reg'].get(type(sock_or_stream))
if not func:
raise Exception('Unknown object for register')
return func(sock_or_stream, events, callback)
def register_socket(self, sock, events, callback):
return self._selector.register(sock, events, callback)
def register_stream(self, stream: TorStream, events, callback):
if events & EVENT_WRITE:
raise Exception('Write event not supported yet')
stream.register(callback)
if stream not in self._stream_to_callback:
self._stream_to_callback[stream] = []
self._stream_to_callback[stream].append(callback)
def unregister(self, sock_or_stream):
func = self._regs_funcs_map['unreg'].get(type(sock_or_stream))
if not func:
raise Exception('Unknown object for unregister')
return func(sock_or_stream)
def unregister_socket(self, sock):
return self._selector.unregister(sock)
def unregister_stream(self, stream):
callbacks = self._stream_to_callback.pop(stream, None)
if not callbacks:
raise Exception('There is no such stream registered')
for callback in callbacks:
stream.unregister(callback)
def _do_stop(self, raw_socket, mask):
self._do_loop = False
def _do_recv(self, raw_socket, mask):
for cell in self._tor_socket.recv_cell_async():
logger.debug('Cell received: %r', cell)
try:
self._handler_mgr.handle(cell)
except BaseException:
logger.exception('Some handle errors')
def run(self):
logger.debug('Starting...')
while self._do_loop:
events = self._selector.select()
for key, mask in events:
callback = key.data
callback(key.fileobj, mask)
self._cleanup()
logger.debug('Stopped...')
class ChatServer:
def __init__(self, port):
self.sockets = {}
self.names = {}
self.socket_server = socket.socket(
socket.AF_INET, socket.SOCK_STREAM | socket.SOCK_NONBLOCK)
self.selector = DefaultSelector()
self.socket_server.bind(("", port))
self._running = False
def _accept(self, socket):
socket_client = socket.accept()[0]
socket_client.setblocking(False)
self.selector.register(socket_client, EVENT_READ, self._read)
def _read(self, socket):
data = socket.recv(256)
data = data.decode('utf-8')
if not data: # socket closed.
name = self.sockets.pop(socket, None)
if name is not None:
self.names.pop(name, None)
return
try:
order, p1, p2 = data.split(';')
except ValueError:
socket.send('ERROR: Bad request {!r}'.format(data).encode("utf-8"))
print('Bad request {!r}'.format(data))
return
if order == "MY_NAME_IS" and socket in self.sockets:
socket.send(b"ERROR: I know you")
return
if order in ('SEND_TO', 'BYE') and socket not in self.sockets:
socket.send(b"ERROR: I don't know you")
return
if order == "MY_NAME_IS":
if p1 in self.names:
socket.send(b'ERROR: nickname is already used')
else:
self.names[p1] = socket
self.sockets[socket] = p1
socket.send(b'OK')
elif order == "SEND_TO":
if p1 not in self.names:
socket.send(b'ERROR: recipient is not connected')
else:
self.names[p1].send(p2.encode('utf-8'))
elif order == "BYE":
name = self.sockets.pop(socket, None)
self.names.pop(name, None)
socket.send(b'BYE')
def run_forever(self):
self.socket_server.listen(1)
self.selector.register(self.socket_server, EVENT_READ, self._accept)
self._running = True
while self._running:
for key, events in self.selector.select():
key.data(key.fileobj)
def close(self):
self._running = False
self.selector.close()
self.socket_server.close()
class DBusConnectionBase:
"""Connection machinery shared by this module and threading"""
def __init__(self, sock: socket.socket, enable_fds=False):
self.sock = sock
self.enable_fds = enable_fds
self.parser = Parser()
self.outgoing_serial = count(start=1)
self.selector = DefaultSelector()
self.select_key = self.selector.register(sock, EVENT_READ)
self.unique_name = None
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
return False
def _serialise(self, message: Message,
serial) -> (bytes, Optional[array.array]):
if serial is None:
serial = next(self.outgoing_serial)
fds = array.array('i') if self.enable_fds else None
data = message.serialise(serial=serial, fds=fds)
return data, fds
def _send_with_fds(self, data, fds):
bytes_sent = self.sock.sendmsg(
[data], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds)])
# If sendmsg succeeds, I think ancillary data has been sent atomically?
# So now we just need to send any leftover normal data.
if bytes_sent < len(data):
self.sock.sendall(data[bytes_sent:])
def _receive(self, deadline):
while True:
msg = self.parser.get_next_message()
if msg is not None:
return msg
b, fds = self._read_some_data(
timeout=deadline_to_timeout(deadline))
self.parser.add_data(b, fds=fds)
def _read_some_data(self, timeout=None):
for key, ev in self.selector.select(timeout):
if key == self.select_key:
if self.enable_fds:
return self._read_with_fds()
else:
return unwrap_read(self.sock.recv(4096)), []
raise TimeoutError
def _read_with_fds(self):
nbytes = self.parser.bytes_desired()
data, ancdata, flags, _ = self.sock.recvmsg(nbytes, fds_buf_size())
if flags & getattr(socket, 'MSG_CTRUNC', 0):
self.close()
raise RuntimeError("Unable to receive all file descriptors")
return unwrap_read(data), FileDescriptor.from_ancdata(ancdata)
def close(self):
"""Close the connection"""
self.selector.close()
self.sock.close()
class Bot:
def __init__(self, name="masshl"):
self.connections = []
self.selector = DefaultSelector()
self.config = Config()
self.running = False
self.is_restarting = False
self.plugins = {}
self.cwd = pathlib.Path().resolve()
self.hooks: Dict[str, List[Hook]] = defaultdict(list)
self.name = name
self.log = Logger(self)
self.storage: DefaultDict[str, Dict] = defaultdict(dict)
def run(self):
self._load_plugins()
self.log.debug(self.hooks)
for network in self.config["connections"]:
temp_connection = Connection(
config=self.config["connections"][network],
selector=self.selector,
bot=self,
name=network,
debug=self.config["debug"])
self.connections.append(temp_connection)
gc.collect()
for connection in self.connections:
connection.connect()
self.running = True
while self.running:
event = self.selector.select(1)
for file, _ in event:
file.fileobj.read()
self.selector.close()
return self.is_restarting
def stop(self, reason=None):
if reason is None:
reason = "Controller requested stop"
for connection in self.connections:
if not connection.hasquit:
connection.quit(reason)
time.sleep(1)
for connection in self.connections:
if connection.connected:
connection.close()
self.running = False
def restart(self, reason=None):
if reason is None:
reason = "Controller requested restart"
self.is_restarting = True
self.stop(reason)
def _load_plugins(self):
path = pathlib.Path("plugins").resolve().relative_to(self.cwd)
for file in path.glob("*.py"):
self.load_plugin('.'.join(file.parts).rsplit('.', 1)[0])
def load_plugin(self, name: Union[List, str]):
if name == "*":
resp = self.load_plugin(list(self.plugins.keys()))
for r in resp:
print(r)
return
if isinstance(name, list):
resp = []
for plugin in name:
print(self.load_plugin(plugin))
return resp
print(f"LOADING {name}")
if not name.startswith("plugins."):
name = "plugins." + name
try:
print("loading plugin", name)
imported_module = importlib.import_module(name)
if hasattr(imported_module, "_masshl_loaded"):
print(f"CALLING RELOAD FOR {name}")
importlib.reload(imported_module)
if name in self.plugins:
print(f"CALLING UNLOAD FOR {name}")
self.unload(name)
else:
setattr(imported_module, "_masshl_loaded", None)
except Exception as e:
self.log.exception(e)
return e
else:
setattr(imported_module, "_masshl_loaded", None)
self._load_hooks(imported_module, name, "on_load*")
responses = self.call_hook(f"on_load_{name}")
ok = True
for resp in responses:
if isinstance(resp, Exception):
ok = False
if not ok:
self.log.error(f"Plugin {name} failed to load.")
return
self.plugins[name] = imported_module
self._load_hooks(imported_module, name)
def unload(self, name):
if not name.startswith("plugins."):
name = "plugins." + name
if name in self.plugins:
del self.plugins[name]
self.log.debug(f"REMOVING PLUGIN: {name}")
todo = []
for hook_name, hook_list in self.hooks.items():
for hook in reversed(hook_list):
if hook.plugin == name:
if hook_name == f"on_unload_{name}":
todo.extend(self.call_hook(f"on_unload_{name}"))
self.log.debug(f"REMOVING HOOK: {hook_name}: {hook}")
hook_list.remove(hook)
self._cleanup_hooks()
self.handle_todos(todo)
def _cleanup_hooks(self):
new_hooks = {n: h for n, h in self.hooks.items() if h}
self.hooks.clear()
self.hooks.update(new_hooks)
print(self.hooks)
def _load_hooks(self, plugin, name, filters: str = None):
self.log.debug(f"Loading {name}'s hooks" +
(f" Filtered to {filters}" if filters else ""))
for func in plugin.__dict__.values():
# self.log.debug(f"checking {func} in {plugin}")
if not hasattr(func, "_IsHook"):
continue
hooks = getattr(func, "_IsHook")
for hook, perm in hooks:
if filters is not None and not fnmatch(hook, filters):
self.log.debug(
f"SKIPPING {hook}: {func}, does not match filter ('{filters}'): '{hook}''"
)
continue
self.log(f"loading new hook {hook}: {func}" +
(f" Hook requested {perm}" if perm else ""))
self.hooks[hook].append(Hook(name, func, perm))
delattr(func, "_IsHook")
def launch_hook_func(self, func: Callable, **kwargs):
sig = inspect.signature(func)
kwargs["bot"] = self
args = []
for arg in sig.parameters:
assert arg in kwargs, \
f"Callback requested an argument that the hook launcher was not passed. it was '{arg}'"
args.append(kwargs[arg])
return func(*args)
def call_hook(self, name, **kwargs):
todos = []
name = name.lower()
if name not in self.hooks:
return todos
for hook in self.hooks[name]:
if hook.perms and not check(kwargs["msg"], hook.perms):
kwargs["msg"].origin.send_notice(
"Sorry, you are not allowed to use this command")
continue
try:
resp = self.launch_hook_func(hook.func, **kwargs)
except Exception as e:
todos.append((hook, e))
self.log(f"Exception in {name}: {hook}")
self.log.exception(e)
else:
todos.append((hook, resp))
return todos
def handle_todos(self, todos, ret=False):
self.log.debug(f"HANDLING TODOS: {todos}")
resp = []
for hook, todo in todos:
if callable(todo):
todo()
elif ret:
resp.append(todo)
else:
self.log(f"{hook}: {todo}")
return resp
class Peer(HandleableMixin, CommandableMixin, LanTopologyMixin,
DefaultAuthenticatorMixin):
"""
Attributes:
server_info (PeerInfo): Peer's peer info.
pkt_handlers (Dict[str, Handler]): All handlers which peer have ability
to process.
peer_pool (Dict[Tuple[str, int], PeerInfo]): All peers currently avai-
lable in net.
"""
@property
def server_info(self):
return self.__server_info
@property
def program_hash(self):
return self.__program_hash
@property
def send_queue(self):
# def packet_queue(self):
return self.__packet_queue
@property
def connectlist(self) -> List: # Remember to remove `List` import.
return self.peer_pool.values()
def __init__(
self,
host: Tuple[str, int],
name: str,
role: str,
cert: Tuple[str, str],
program_hash: str,
ns: str,
auto_register: bool = False,
logger: "logging.Logger" = getLogger(__name__),
) -> None:
"""Init of PeerManager
Args:
host: Binding host.
name: Peer's name in net.
role: Peer's role in net.
cert: Cert file's path.
program_hash: Program self hash to send in packet.
ns: Nameserver address for resolve DNS.
logger: Logger for logging.
"""
super().__init__()
self.logger = getLogger(name)
self.__auto_register = auto_register
self.__selector = DefaultSelector()
self.__packet_queue = {}
self.__cert = cert
self.__program_hash = program_hash
self.__server_info = PeerInfo(host=host, name=name, role=role)
self.__tcp_server = self.__bind_socket(cert=self.__cert)
self.peer_pool = {}
self.pkt_handlers = {}
self.commands = {}
self.logger.info("Program hash: {{{}...{}}}".format(
self.__program_hash[:6], self.__program_hash[-6:]))
self.dns_resolver = DNSResolver(ns="127.0.0.1" if ns is None else ns,
role=role)
self.monitor = Monitor(peer=self, logger=getLogger(name + ".MONITOR"))
if self.__auto_register is False:
self.logger.warning(
("auto_register parameter is set to False,\n You may need to r"
"egister them through _register_command & _register_handler m"
"ethod."))
def _preregister_handler(self) -> None:
self.topology_register_handler()
installing_handlers = [MessageHandler(self)]
for each in installing_handlers:
self.register_handler(handler=each)
def _preregister_command(self) -> None:
installing_commands = [
HelpCmd(self),
JoinCmd(self),
SendCmd(self),
ListCmd(self),
LeaveNetCmd(self),
]
for each in installing_commands:
self.register_command(command=each)
def pend_packet(self, sock: "Socket", pkt: "Packet", **kwargs) -> None:
"""Pending pkt's raw_data to queue's with sepecific sock.
Any exception when wrapping handler to packet whould cause this connec-
tion been close and thread maintaining loop terminate.
Args:
sock: A Socket which wants to pend on its queue.
pkt: A Packet ready to be pend.
**kwargs: Any additional arguments needs by handler object.
Raises:
AssertionError:
If given pkt variable is not in proper Packet type.
"""
assert type(pkt) is Packet
try:
self.__packet_queue[sock].put_nowait(pkt)
except Exception:
self.logger.info(format_exc())
def register_socket(self, sock: "Socket") -> None:
"""Register a new socket with packet queue & selector.
Init a packet queue and put into dict for further handling of packets.
And the given socket will be register in selector for IO process.
Args:
sock: A Socket object which wants to be register.
"""
self.__packet_queue[sock] = Queue()
self.__selector.register(sock, EVENT_READ | EVENT_WRITE,
self.__on_handle)
def unregister_socket(self, sock: "Socket") -> None:
del self.__packet_queue[sock]
self.__selector.unregister(sock)
def _on_packet(self, sock: "Socket", pkt: "Packet",
handler: "Handler") -> None:
"""Method use to process passed packet to higher application layer.
This method will call by AuthenticatorMixin when a packet is passed examination.
This is 3rd layer to process packet to handler. This is last layer to application layer.
"""
handler.on_recv(src=pkt.src, pkt=pkt, sock=sock)
self.monitor.on_recv_pkt(addr=pkt.src, pkt=pkt, conn=sock)
def new_tcp_long_conn(self, dst: Tuple[str, int]) -> "SSLSocket":
"""Create a ssl-wrapped TCP socket with given destination host
Args:
dst: Specified socket destination.
Returns:
A SSLSocket object which connected to destination host with
non-blocking.
Raises:
AssertionError:
If given dst variable is not in proper Tuple[str, int] type.
"""
assert host_valid(dst) is True
unwrap_socket = socket(AF_INET, SOCK_STREAM)
sock = wrap_socket(unwrap_socket,
cert_reqs=CERT_REQUIRED,
ca_certs=self.__cert[0])
sock.connect(dst)
sock.setblocking(False)
return sock
def loop_start(self):
self.logger.info(self.__server_info)
self.__selector.register(self.__tcp_server, EVENT_READ,
self.__on_accept)
if self.__auto_register is True:
self._preregister_handler()
self._preregister_command()
if self.monitor.is_start() is False:
self.monitor.start()
self.logger.info("Peer started.")
def loop(self):
return self._loop()
def _loop(self):
"""Called inside infinite loop from outside inherited class.
It's use to call the method which given event is triggered.
"""
events = self.__selector.select(timeout=0)
for key, mask in events:
if callable(key.data):
key.data(key.fileobj, mask)
def loop_stop(self):
for _, value in self.__packet_queue.items():
if value.empty() is False:
sleep(2)
return self.loop_stop()
self.__selector.unregister(self.__tcp_server)
self.leave_net()
self.monitor.stop()
def loop_stop_post(self):
self.__tcp_server.close()
self.__selector.close()
def __bind_socket(self, cert: Tuple[str, str]) -> "SSLSocket":
unwrap_socket = socket(AF_INET, SOCK_STREAM)
unwrap_socket.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
unwrap_socket.setsockopt(SOL_SOCKET, SO_REUSEPORT, 1)
unwrap_socket.bind(self.__server_info.host)
unwrap_socket.listen(5)
unwrap_socket.setblocking(False)
self.logger.info("Peer prepared")
self.logger.info(
"This peer is running with certificate at path {}".format(cert[0]))
self.logger.info("Please make sure other peers have same certicate.")
return wrap_socket(unwrap_socket,
certfile=cert[0],
keyfile=cert[1],
server_side=True)
def __on_accept(self, sock: "Socket", mask, **kwargs):
"""Call when a new socket is connection by waiter socket.
This will accept all sockets from outside, but it doesn't mean every socket's
packet be process by higher layer.
This is the 1st layer to process sockets.
"""
conn, _ = sock.accept()
conn.setblocking(False)
self.register_socket(sock=conn)
def __on_handle(self, sock: "Socket", mask, **kwargs):
"""Decide whether send or recv."""
if mask & EVENT_READ == EVENT_READ:
self.__on_recv(sock=sock, mask=mask, **kwargs)
if mask & EVENT_WRITE == EVENT_WRITE:
self.__on_send(sock=sock, mask=mask, **kwargs)
def __on_recv(self, sock: "Socket", mask, **kwargs):
"""Method use when recieve socket data.
This is 2th layer to process Packets.
"""
try:
raw_data = sock.recv(4096)
if raw_data == b"":
return
pkt = Packet.deserilize(raw_data=raw_data)
return self._authenticate_packet(sock=sock, pkt=pkt)
except SSLWantReadError:
return
except sock_error as sock_err:
if sock_err.errno == ECONNRESET:
peer_info = self.get_peer_info_by_conn(conn=sock)
if peer_info is not None:
peer_info.status.update(status_type=StatusType.NO_RESP)
self.logger.warning("Peer {} Connection Reseted.".format(
peer_info.host))
else:
raise sock_err
except Exception:
self.logger.warning(str(self.server_info) + format_exc())
def __on_send(self, sock: "Socket", mask, **kwargs):
"""Method use when sending data to socket."""
q = self.__packet_queue[sock] if sock in self.__packet_queue else None
while q is not None and q.empty() is False:
try:
pkt = q.get_nowait()
handler = self.select_handler(pkt_type=pkt._type)
handler.pre_send(pkt=pkt)
data = Packet.serilize(obj=pkt)
sock.send(data)
handler.post_send(pkt=pkt, sock=sock)
except sock_error as sock_err:
if sock_err.errno == ECONNRESET:
q.put_nowait(pkt)
self.monitor.peer_status_update_by_host(
host=pkt.src, status_type=StatusType.NO_RESP)
self.logger.warning("Peer {} Connection Reseted.".format(
pkt.src))
else:
raise sock_err
except Exception:
self.logger.warning(format_exc())
def distribute(cmd, max_bytes, max_procs, chunk_size, round_robin, verbose):
"""
Blocking function that manages all the delegation of chunks from
stdin to subprocesses, spawning of subprocess, and collation of
stdout and stderr from each subprocess.
Broadly speaking, gatling reads chunks of data from stdin and
disperses those among subprocesses to achieve parallel execution.
Stdin is read in chunks of chunk_size bytes and truncated to the
last newline, keeping the remainder to be prepended on the next
chunk. Multiple chunks without newline is allowed, nothing is
passed on until a newline is found or stdin closes.
The stdout and stderr from each child is read in a similar
fashion. Whenever a newline is found the output is written onto
stdout or stderr, respectively. This collation preserves the
format of the output, but only weakly adheres to the chronology.
The output from gatling is the result of the subprocesses with a
line-by-line integrity preserved with no guarantee that the order
is exactly maintained.
There are two different behaviors for subprocess spawning,
manifold or gatling:
* In manifold-mode each new chunk read from stdin spawns a new
subprocess until max_procs is reached and then each of the
subprocesses are fed a chunk in round robin fashion. This is an
excellent model for programs that do not tax an external
resource and programs that can act on stdin as soon as it is
available.
* In gatling-mode chunks are fed to a single subprocess until that
processes' max_bytes is reached, the subprocess' stdin is closed
and on the next chunk a new subprocess is spawned and fed until
its max_bytes is reached and so on until max_procs is reached.
If max_procs is reached gatling is blocked until a subprocess
finishes. This mode works well for programs that connect to an
external service or programs that don't start processing stdin
until it is closed.
cmd (list): The full command line for spawning subprocesses.
max_bytes (int): Maximum bytes to pass to each subprocess before
closing the subprocess' stdin. Increase this value if
subprocesses do not have memory management problems from large
input and if subprocesses process stdin as stream. Decrease
this value if programs can not handle a large input set on
stdin or if programs do not start processing until stdin is
closed.
max_procs (int): Maximum number of simultaneos subprocesses.
Increase this number if the subprocess is largely CPU bound,
decrease it to match the hardware if the subprocesses are IO
bound.
chunk_size (int): Stdin content streamed to gatling is consumed in
chunks of this size (bytes) for efficiency reasons. Originally
it was line by line, but that was too slow to keep
subprocesses fed continuously. To force line-by-line behavior,
set chunk to a size always less than the length of an input
line (worst case 1, but try to keep it as high as possible).
Experiment with this number to maximize pipeline throughput.
round_robin (bool): True is manifold, false is gatling, see above
for details.
"""
from sys import stdin, stdout, stderr
from subprocess import Popen, PIPE
from time import sleep, time
from selectors import DefaultSelector, EVENT_READ
from threading import Thread
from fcntl import fcntl, F_SETFL, F_GETFL
from os import O_NONBLOCK
def sink(selector, not_done):
n = 0
while len(not_done) > 1 or list(selector.get_map()):
for (fileobj, _, _, p), _ in selector.select():
chunk = fileobj.read(4096)
if chunk:
i = chunk.rfind('\n') + 1
if i:
n += fileobj._trg.write(fileobj._buf)
n += fileobj._trg.write(chunk[:i])
fileobj._buf = chunk[i:]
else:
fileobj._buf += chunk
else:
if p.returncode is not None:
n += fileobj._trg.write(fileobj._buf)
selector.unregister(fileobj)
not_done[0] = n
my_name = 'manifold' if round_robin else 'gatling'
selector = DefaultSelector()
res = []
p_filled = [] # Child processes that have had their maximum input supplied
# (p_open) Child processes that can take more input
# In non-round-robin mode, there will only be one such child,
# which will continually be popped and re-appended.
# In round-robin mode, this list is rotated each time there is
# a buffer to be written to a child.
p_open = []
b_in = 0
buf = ''
not_done = [0, 1]
sel_t = None
t0 = time()
try:
for chunk in iter(lambda: stdin.read(chunk_size), ''):
b_in += len(chunk)
i = chunk.rfind('\n') + 1
if i:
p = None
if round_robin:
if len(p_open) + len(p_filled) == max_procs:
p = p_open.pop(0)
else:
if p_open:
p = p_open.pop(0)
if not p:
if verbose:
running = len(p_filled) + len(p_open)
print(
f"# {my_name} STARTED A PROCESS (1 + {running} + {len(res)}):",
*cmd,
file=sys.stderr)
p = Popen(cmd,
encoding=stdout.encoding,
stdin=PIPE,
stdout=PIPE,
stderr=PIPE)
p._n = 0
p._t0 = time()
for fo, trg in [(p.stdout, stdout), (p.stderr, stderr)]:
fo._buf = ''
fo._trg = trg
fcntl(fo, F_SETFL, fcntl(fo, F_GETFL) | O_NONBLOCK)
selector.register(fo, EVENT_READ, p)
if not sel_t:
sel_t = Thread(target=sink, args=(selector, not_done))
sel_t.daemon = True
sel_t.start()
p._n += p.stdin.write(buf)
p._n += p.stdin.write(chunk[:i])
buf = chunk[i:]
if p._n >= max_bytes:
t1 = time()
p._t1 = t1
td = (t1 - t0) * 1024
ptd = (t1 - p._t0) * 1024
p.stdin.close()
p_filled.append(p)
if verbose and ptd and td:
running = len(p_filled) + len(p_open)
print(
f"# {my_name} PROCESS LIMIT {p._n:,}/{max_bytes:,}",
f"({p._n/ptd:.2f} kb/s).",
f"INPUT: {b_in:,} ({b_in/td:.2f} kb/s) OUTPUT: {not_done[0]:,}.",
f"PROCESSES: {running}/{running+len(res)}",
file=sys.stderr)
while len(p_filled) == max_procs:
done = [d for d in p_filled if d.poll() is not None]
if done and verbose:
print(f"# {my_name} CLOSED {len(done)} PROCESSES",
file=sys.stderr)
for d in done:
if verbose and p._t0 != t1:
print(
f"# {my_name} CLOSED PROCESS INPUT: {p._n:,} TIME:",
f"{t1-p._t0:.1f}/{t1-p._t1:.1f}",
f"KB/S: {p._n/(t1-p._t0)/1024:.2f}",
file=sys.stderr)
p_filled.remove(d)
res.append(d.returncode)
if not done:
sleep(0.5)
else:
p_open.append(p)
else:
buf += chunk
for p in p_open:
p.stdin.write(buf)
buf = ''
p.stdin.close()
p_filled.append(p)
except (KeyboardInterrupt, SystemExit):
not_done.pop()
for p in p_open:
p.stdin.close()
p.kill()
for d in p_filled:
d.kill()
raise
while p_filled:
res.append(p_filled.pop(0).wait())
if sel_t:
not_done.pop()
sel_t.join()
selector.close()
return res
class DBusConnection:
def __init__(self, sock):
self.sock = sock
self.parser = Parser()
self.outgoing_serial = count(start=1)
self.selector = DefaultSelector()
self.select_key = self.selector.register(sock, EVENT_READ)
self._unwrap_reply = False
# Message routing machinery
self.router = Router(_Future) # Old interface, for backwards compat
self._filters = MessageFilters()
# Say Hello, get our unique name
self.bus_proxy = Proxy(message_bus, self)
hello_reply = self.bus_proxy.Hello()
self.unique_name = hello_reply[0]
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
return False
def send(self, message: Message, serial=None):
"""Serialise and send a :class:`~.Message` object"""
if serial is None:
serial = next(self.outgoing_serial)
data = message.serialise(serial=serial)
self.sock.sendall(data)
send_message = send # Backwards compatibility
def receive(self, *, timeout=None) -> Message:
"""Return the next available message from the connection
If the data is ready, this will return immediately, even if timeout<=0.
Otherwise, it will wait for up to timeout seconds, or indefinitely if
timeout is None. If no message comes in time, it raises TimeoutError.
"""
deadline = timeout_to_deadline(timeout)
while True:
msg = self.parser.get_next_message()
if msg is not None:
return msg
b = self._read_some_data(timeout=deadline_to_timeout(deadline))
self.parser.add_data(b)
def _read_some_data(self, timeout=None):
for key, ev in self.selector.select(timeout):
if key == self.select_key:
return unwrap_read(self.sock.recv(4096))
raise TimeoutError
def recv_messages(self, *, timeout=None):
"""Receive one message and apply filters
See :meth:`filter`. Returns nothing.
"""
msg = self.receive(timeout=timeout)
self.router.incoming(msg)
for filter in self._filters.matches(msg):
filter.queue.append(msg)
def send_and_get_reply(self, message, *, timeout=None, unwrap=None):
"""Send a message, wait for the reply and return it
Filters are applied to other messages received before the reply -
see :meth:`add_filter`.
"""
check_replyable(message)
deadline = timeout_to_deadline(timeout)
if unwrap is None:
unwrap = self._unwrap_reply
serial = next(self.outgoing_serial)
self.send_message(message, serial=serial)
while True:
msg_in = self.receive(timeout=deadline_to_timeout(deadline))
reply_to = msg_in.header.fields.get(HeaderFields.reply_serial, -1)
if reply_to == serial:
if unwrap:
return unwrap_msg(msg_in)
return msg_in
# Not the reply
self.router.incoming(msg_in)
for filter in self._filters.matches(msg_in):
filter.queue.append(msg_in)
def filter(self, rule, *, queue: Optional[deque] = None, bufsize=1):
"""Create a filter for incoming messages
Usage::
with conn.filter(rule) as matches:
# matches is a deque containing matched messages
matching_msg = conn.recv_until_filtered(matches)
:param jeepney.MatchRule rule: Catch messages matching this rule
:param collections.deque queue: Matched messages will be added to this
:param int bufsize: If no deque is passed in, create one with this size
"""
return FilterHandle(self._filters, rule, queue
or deque(maxlen=bufsize))
def recv_until_filtered(self, queue, *, timeout=None) -> Message:
"""Process incoming messages until one is filtered into queue
Pops the message from queue and returns it, or raises TimeoutError if
the optional timeout expires. Without a timeout, this is equivalent to::
while len(queue) == 0:
conn.recv_messages()
return queue.popleft()
In the other I/O modules, there is no need for this, because messages
are placed in queues by a separate task.
:param collections.deque queue: A deque connected by :meth:`filter`
:param float timeout: Maximum time to wait in seconds
"""
deadline = timeout_to_deadline(timeout)
while len(queue) == 0:
self.recv_messages(timeout=deadline_to_timeout(deadline))
return queue.popleft()
def close(self):
"""Close this connection"""
self.selector.close()
self.sock.close()
