QuLab 2.10.10__cp313-cp313-macosx_10_13_universal2.whl

qulab/sys/net/bencoder.py

@@ -0,0 +1,205 @@
+"""
+This module provides a simple bencoding implementation in pure Python.
+Bencoding is a way to specify and organize data in a terse format. It supports
+the following types: byte strings, integers, lists, and dictionaries.
+
+Classes:
+    Bencached: A class that holds bencoded data.
+
+Functions:
+    encode_bencached(x, fp): Encodes a Bencached object.
+    encode_int(x, fp): Encodes an integer.
+    encode_bool(x, fp): Encodes a boolean.
+    encode_string(x, fp): Encodes a string.
+    encode_list(x, fp): Encodes a list.
+    encode_dict(x, fp): Encodes a dictionary.
+    decode_int(x, ptr): Decodes a bencoded integer.
+    decode_string(x, ptr): Decodes a bencoded string.
+    decode_list(x, ptr): Decodes a bencoded list.
+    decode_dict(x, ptr): Decodes a bencoded dictionary.
+    encode(obj, fp): Encodes an object.
+    encodes(obj): Encodes an object and returns the bencoded bytes.
+    decode(fp): Decodes a file-like object.
+    decodes(s): Decodes a bencoded bytes object.
+    dump = encode: Alias for encode function.
+    load = decodes: Alias for decode function.
+    dumps = encodes: Alias for encodes function.
+    loads = decodes: Alias for decodes function.
+
+Examples:
+    >>> encode(42) == b'i42e'
+    True
+    >>> decode(b'i42e')
+    42
+"""
+import io
+
+
+class Bencached():
+
+    __slots__ = ['bencoded']
+
+    def __init__(self, s):
+        self.bencoded = s
+
+
+def encode_bencached(x, fp):
+    fp.write(x.bencoded)
+
+
+def encode_int(x, fp):
+    fp.write(f"i{x}e".encode())
+
+
+def encode_bool(x, fp):
+    if x:
+        encode_int(1, fp)
+    else:
+        encode_int(0, fp)
+
+
+def encode_string(x, fp):
+    if isinstance(x, str):
+        x = x.encode()
+    fp.write(f"{len(x):d}:".encode())
+    fp.write(x)
+
+
+def encode_list(x, fp):
+    fp.write(b'l')
+    for i in x:
+        encode_func[type(i)](i, fp)
+    fp.write(b'e')
+
+
+def encode_dict(x, fp):
+    fp.write(b'd')
+    for k, v in sorted(x.items()):
+        if isinstance(k, str):
+            k = k.encode()
+        fp.write(f"{len(k):d}:".encode())
+        fp.write(k)
+        encode_func[type(v)](v, fp)
+    fp.write(b'e')
+
+
+encode_func = {
+    Bencached: encode_bencached,
+    int: encode_int,
+    str: encode_string,
+    bytes: encode_string,
+    list: encode_list,
+    tuple: encode_list,
+    dict: encode_dict,
+    bool: encode_bool
+}
+
+
+def decode_int(x, ptr):
+    ptr += 1
+    end = x.index(b'e', ptr)
+    if x[ptr] == b'-'[0]:
+        if x[ptr + 1] == b'0'[0]:
+            raise ValueError
+    elif x[ptr] == b'0'[0] and end != ptr + 1:
+        raise ValueError
+    n = int(x[ptr:end])
+    return n, end + 1
+
+
+def decode_string(x, ptr):
+    colon = x.index(b':', ptr)
+    n = int(x[ptr:colon])
+    if x[ptr] == b'0'[0] and colon != ptr + 1:
+        raise ValueError
+    colon += 1
+    return x[colon:colon + n], colon + n
+
+
+def decode_list(x, ptr):
+    r, ptr = [], ptr + 1
+    while x[ptr] != b'e'[0]:
+        v, ptr = decode_func[x[ptr]](x, ptr)
+        r.append(v)
+    return r, ptr + 1
+
+
+def decode_dict(x, ptr):
+    r, ptr = {}, ptr + 1
+    while x[ptr] != b'e'[0]:
+        k, ptr = decode_string(x, ptr)
+        if isinstance(k, bytes):
+            k = k.decode()
+        r[k], ptr = decode_func[x[ptr]](x, ptr)
+    return r, ptr + 1
+
+
+decode_func = {b'l'[0]: decode_list, b'd'[0]: decode_dict, b'i'[0]: decode_int}
+
+for i in range(10):
+    decode_func[f"{i}".encode()[0]] = decode_string
+
+
+def encode(obj, fp):
+    try:
+        encode_func[type(obj)](obj, fp)
+    except KeyError:
+        raise ValueError("Allowed types: int, bytes, list, dict; not %s",
+                         type(obj))
+
+
+def encodes(obj):
+    """
+    bencodes given object. Given object should be a int,
+    bytes, list or dict. If a str is given, it'll be
+    encoded as ASCII.
+    >>> [encode(i) for i in (-2, 42, b"answer", b"")] \
+            == [b'i-2e', b'i42e', b'6:answer', b'0:']
+    True
+    >>> encode([b'a', 42, [13, 14]]) == b'l1:ai42eli13ei14eee'
+    True
+    >>> encode({'bar': b'spam', 'foo': 42, 'mess': [1, b'c']}) \
+            == b'd3:bar4:spam3:fooi42e4:messli1e1:cee'
+    True
+    """
+    buff = io.BytesIO()
+    encode(obj, buff)
+    return buff.getvalue()
+
+
+def decodes(s):
+    """
+    Decodes given bencoded bytes object.
+    >>> decode(b'i-42e')
+    -42
+    >>> decode(b'4:utku') == b'utku'
+    True
+    >>> decode(b'li1eli2eli3eeee')
+    [1, [2, [3]]]
+    >>> decode(b'd3:bar4:spam3:fooi42ee') == {'bar': b'spam', 'foo': 42}
+    True
+    """
+    if isinstance(s, str):
+        s = s.encode()
+    try:
+        r, l = decode_func[s[0]](s, 0)
+    except (IndexError, KeyError, ValueError):
+        raise ValueError("not a valid bencoded string")
+    if l != len(s):
+        raise ValueError("invalid bencoded value (data after valid prefix)")
+    return r
+
+
+def decode(fp):
+    """
+    Decodes given file-like object.
+    """
+    return decodes(fp.read())
+
+
+dump = encode
+load = decodes
+dumps = encodes
+loads = decodes
+
+__all__ = ['decodes', 'encodes', 'dump', 'load', 'dumps', 'loads']

qulab/sys/net/cli.py

@@ -0,0 +1,169 @@
+import asyncio
+import pathlib
+import pickle
+
+import click
+
+from .kad import Server
+
+
+async def client_get(key, server, port, timeout=5):
+    reader, writer = await asyncio.wait_for(asyncio.open_connection(
+        server, port),
+                                            timeout=timeout)
+    writer.write(pickle.dumps(('get', key)))
+    await asyncio.wait_for(writer.drain(), timeout=timeout)
+    data = await asyncio.wait_for(reader.read(1600), timeout=timeout)
+    value = pickle.loads(data)
+    writer.close()
+    await writer.wait_closed()
+    return value
+
+
+async def client_set(key, value, server, port, timeout=5):
+    reader, writer = await asyncio.wait_for(asyncio.open_connection(
+        server, port),
+                                            timeout=timeout)
+    writer.write(pickle.dumps(('set', key, value)))
+    await asyncio.wait_for(writer.drain(), timeout=timeout)
+    writer.close()
+    await writer.wait_closed()
+
+
+async def client_bootstrap(address, server, port, timeout=5):
+    reader, writer = await asyncio.wait_for(asyncio.open_connection(
+        server, port),
+                                            timeout=timeout)
+    addr, p = address.split(':')
+    writer.write(pickle.dumps(('bootstrap', (addr, int(p)))))
+    await asyncio.wait_for(writer.drain(), timeout=timeout)
+    writer.close()
+    await writer.wait_closed()
+
+
+@click.group()
+def dht():
+    pass
+
+
+@dht.command()
+@click.option('--interface',
+              default='127.0.0.1',
+              help='Interface to listen on')
+@click.option('--port', default=8467, help='Port to listen on')
+@click.option('--dht-only', default=False, is_flag=True, help='DHT only')
+@click.option('--dht-interface',
+              default='0.0.0.0',
+              help='DHT interface to listen on')
+@click.option('--dht-port', default=8468, help='DHT port to listen on')
+@click.option('--bootstrap', default=None, help='Address of bootstrap node')
+@click.option('--interval', default=300, help='Refresh interval')
+@click.option('--state-file', default=None, help='State file')
+def server(interface, port, dht_only, dht_interface, dht_port, bootstrap,
+           interval, state_file):
+
+    async def main(interface, port, dht_only, dht_interface, dht_port,
+                   bootstrap, state_file):
+
+        if state_file is None:
+            state_file = pathlib.Path.home() / '.waveforms' / 'dht.pickle'
+        else:
+            state_file = pathlib.Path(state_file)
+
+        if state_file.exists():
+            node = await Server.load_state(state_file, dht_port, dht_interface,
+                                           interval)
+        else:
+            node = Server()
+            state_file.parent.mkdir(parents=True, exist_ok=True)
+            await node.listen(dht_port, dht_interface, interval)
+        if bootstrap:
+            addr, p = bootstrap.split(':')
+            await node.bootstrap([(addr, int(p))])
+
+        loop = asyncio.get_running_loop()
+        loop.call_later(min(interval / 2, 30), node.save_state_regularly,
+                        state_file, interval)
+
+        async def handle_cmds(reader, writer):
+            data = await reader.read(1600)
+            cmd, *args = pickle.loads(data)
+            match cmd:
+                case 'get':
+                    value = await node.get(*args)
+                    data = pickle.dumps(value)
+                case 'set':
+                    value = await node.set(*args)
+                    data = pickle.dumps(value)
+                case 'ping':
+                    data = pickle.dumps(node.node.id)
+                case 'bootstrap':
+                    await node.bootstrap(args)
+                    data = pickle.dumps(None)
+                case _:
+                    data = pickle.dumps(None)
+
+            writer.write(data)
+            await writer.drain()
+            writer.close()
+            await writer.wait_closed()
+
+        if dht_only:
+            try:
+                while True:
+                    await asyncio.sleep(1)
+            except asyncio.CancelledError:
+                pass
+        else:
+            server = await asyncio.start_server(handle_cmds, interface, port)
+
+            async with server:
+                try:
+                    await server.serve_forever()
+                except asyncio.CancelledError:
+                    pass
+
+    asyncio.run(
+        main(interface, port, dht_only, dht_interface, dht_port, bootstrap,
+             state_file))
+
+
+@dht.command('get')
+@click.argument('key')
+@click.option('--server', default='127.0.0.1', help='Server to connect to')
+@click.option('--port', default=8467, help='Port to connect to')
+@click.option('--timeout', default=5, help='Timeout')
+def get(key, server, port, timeout):
+
+    async def client():
+        value = await client_get(key, server, port, timeout)
+        print(value)
+
+    asyncio.run(client())
+
+
+@dht.command('set')
+@click.argument('key')
+@click.argument('value')
+@click.option('--server', default='127.0.0.1', help='Server to connect to')
+@click.option('--port', default=8467, help='Port to connect to')
+@click.option('--timeout', default=5, help='Timeout')
+def set_(key, value, server, port, timeout):
+
+    async def client():
+        await client_set(key, value, server, port, timeout)
+
+    asyncio.run(client())
+
+
+@dht.command('bootstrap')
+@click.argument('address')
+@click.option('--server', default='127.0.0.1', help='Server to connect to')
+@click.option('--port', default=8467, help='Port to connect to')
+@click.option('--timeout', default=5, help='Timeout')
+def bootstrap(address, server, port, timeout):
+
+    async def client():
+        await client_bootstrap(address, server, port, timeout)
+
+    asyncio.run(client())