mirror of
https://github.com/savoirfairelinux/opendht.git
synced 2026-11-08 12:28:42 +00:00
282 lines
8.7 KiB
Python
282 lines
8.7 KiB
Python
# Copyright (c) 2014-2026 Savoir-faire Linux Inc.
|
|
# SPDX-License-Identifier: MIT
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any, Callable, Iterable, List, Optional, Tuple
|
|
|
|
from . import _core as _dht
|
|
|
|
# Re-export core types for convenience
|
|
from ._core import InfoHash, Value, ValueType, Where, SockAddr, PublicKey, PkId, DhtConfig, Identity
|
|
|
|
|
|
class _Listener:
|
|
"""Async iterator around DHT listen callback.
|
|
|
|
Yields tuples (Value, expired: bool).
|
|
"""
|
|
|
|
__slots__ = ("_runner", "_token", "_loop", "_queue", "_closed")
|
|
|
|
def __init__(self, runner: DhtRunner, key: InfoHash) -> None:
|
|
self._runner = runner
|
|
self._loop = runner._loop
|
|
self._queue: asyncio.Queue[Tuple[Value, bool]] = asyncio.Queue()
|
|
self._closed = False
|
|
|
|
def _on_value(v: Value, expired: bool) -> bool:
|
|
if self._closed:
|
|
return False
|
|
self._loop.call_soon_threadsafe(self._queue.put_nowait, (v, expired))
|
|
return True # continue listening
|
|
|
|
self._token = self._runner._dht.listen(key, _on_value)
|
|
|
|
def __enter__(self) -> _Listener:
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
self.cancel()
|
|
|
|
def __aiter__(self) -> _Listener:
|
|
return self
|
|
|
|
async def __anext__(self) -> Tuple[Value, bool]:
|
|
if self._closed:
|
|
raise StopAsyncIteration
|
|
item = await self._queue.get()
|
|
if item is None:
|
|
self._closed = True
|
|
raise StopAsyncIteration
|
|
return item
|
|
|
|
def cancel(self) -> None:
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
try:
|
|
self._runner._dht.cancelListen(self._token)
|
|
finally:
|
|
# Unblock any pending __anext__
|
|
self._queue.put_nowait(None)
|
|
|
|
|
|
class _GetStream:
|
|
"""Async iterator around DHT get callback.
|
|
|
|
Yields Value instances as they are found, then ends.
|
|
Raises on failure.
|
|
"""
|
|
__slots__ = ("_runner", "_loop", "_queue", "_closed", "_error")
|
|
|
|
def __init__(
|
|
self,
|
|
runner: DhtRunner,
|
|
key: InfoHash,
|
|
where: Optional[Where],
|
|
filter: Optional[Callable[[Value], bool]],
|
|
) -> None:
|
|
self._runner = runner
|
|
self._loop = runner._loop
|
|
self._queue: asyncio.Queue[Optional[Value]] = asyncio.Queue()
|
|
self._closed = False
|
|
self._error: Optional[BaseException] = None
|
|
|
|
def _on_value(v: Value) -> bool:
|
|
# Called on DHT thread
|
|
if self._closed:
|
|
return False
|
|
self._loop.call_soon_threadsafe(self._queue.put_nowait, v)
|
|
return True
|
|
|
|
def _on_done(ok: bool, _nodes: Iterable[Any]) -> None:
|
|
# Called on DHT thread
|
|
if self._closed:
|
|
return
|
|
if not ok:
|
|
self._error = RuntimeError("DHT get failed")
|
|
self._loop.call_soon_threadsafe(self._queue.put_nowait, None)
|
|
|
|
self._runner._dht.get(key, _on_value, _on_done, filter, where)
|
|
|
|
def __enter__(self) -> _GetStream:
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb) -> None:
|
|
self.cancel()
|
|
|
|
def __aiter__(self) -> _GetStream:
|
|
return self
|
|
|
|
async def __anext__(self) -> Value:
|
|
if self._closed:
|
|
raise StopAsyncIteration
|
|
item = await self._queue.get()
|
|
if item is None:
|
|
self._closed = True
|
|
if self._error is not None:
|
|
raise self._error
|
|
raise StopAsyncIteration
|
|
return item
|
|
|
|
def cancel(self) -> None:
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
# Unblock any pending __anext__
|
|
self._queue.put_nowait(None)
|
|
|
|
|
|
class DhtRunner:
|
|
"""Async wrapper around opendht.DhtRunner.
|
|
"""
|
|
|
|
__slots__ = ("_dht", "_loop", "_config", "_bootstrap")
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
loop: Optional[asyncio.AbstractEventLoop] = None,
|
|
core: Optional[_dht.DhtRunner] = None,
|
|
config: Optional[DhtConfig] = None,
|
|
bootstrap: Optional[Iterable[Tuple[str, Optional[str]]]] = None,
|
|
) -> None:
|
|
try:
|
|
self._loop = loop or asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
self._loop = asyncio.get_event_loop()
|
|
self._dht = core or _dht.DhtRunner()
|
|
self._config = config
|
|
self._bootstrap = list(bootstrap) if bootstrap is not None else None
|
|
|
|
def run(
|
|
self,
|
|
id: Optional[Identity] = None,
|
|
is_bootstrap: bool = False,
|
|
port: int = 0,
|
|
ipv4: str = "",
|
|
ipv6: str = "",
|
|
config: Optional[DhtConfig] = None,
|
|
) -> None:
|
|
cfg = config or self._config or DhtConfig()
|
|
self._dht.run(id, is_bootstrap, port, ipv4, ipv6, cfg)
|
|
if self._bootstrap:
|
|
for host, port_str in self._bootstrap:
|
|
self._dht.bootstrap(host, port_str)
|
|
|
|
def bootstrap(self, host: str, port: Optional[str] = None) -> None:
|
|
self._dht.bootstrap(host, port)
|
|
|
|
async def ping(self, addr: SockAddr) -> bool:
|
|
fut: asyncio.Future[bool] = self._loop.create_future()
|
|
self._dht.ping(addr, lambda ok: self._loop.call_soon_threadsafe(fut.set_result, bool(ok)))
|
|
return await fut
|
|
|
|
def get(
|
|
self,
|
|
key: InfoHash,
|
|
*,
|
|
where: Optional[Where] = None,
|
|
filter: Optional[Callable[[Value], bool]] = None,
|
|
):
|
|
# Return a streaming iterator of values; it ends on completion and raises on failure.
|
|
return _GetStream(self, key, where, filter)
|
|
|
|
async def getAll(
|
|
self,
|
|
key: InfoHash,
|
|
*,
|
|
where: Optional[Where] = None,
|
|
filter: Optional[Callable[[Value], bool]] = None,
|
|
) -> List[Value]:
|
|
"""Collect all values for a key and return them as a list."""
|
|
results: List[Value] = []
|
|
async for v in self.get(key, where=where, filter=filter):
|
|
results.append(v)
|
|
return results
|
|
|
|
async def getFirst(
|
|
self,
|
|
key: InfoHash,
|
|
*,
|
|
where: Optional[Where] = None,
|
|
predicate: Optional[Callable[[Value], Any]] = None,
|
|
filter: Optional[Callable[[Value], bool]] = None,
|
|
) -> Optional[Value]:
|
|
"""Return the first value matching an (optionally async) predicate, or None."""
|
|
with self.get(key, where=where, filter=filter) as stream:
|
|
async for value in stream:
|
|
if predicate is None:
|
|
return value
|
|
result = predicate(value)
|
|
if asyncio.iscoroutine(result):
|
|
result = await result
|
|
if result:
|
|
return value
|
|
return None
|
|
|
|
async def put(self, key: InfoHash, val: Value, *, permanent: bool = False) -> bool:
|
|
fut: asyncio.Future[bool] = self._loop.create_future()
|
|
|
|
def _done(ok: bool, _nodes: Iterable[Any]) -> None:
|
|
self._loop.call_soon_threadsafe(fut.set_result, bool(ok))
|
|
|
|
self._dht.put(key, val, _done, permanent)
|
|
return await fut
|
|
|
|
async def putSigned(self, key: InfoHash, val: Value, *, permanent: bool = False) -> bool:
|
|
fut: asyncio.Future[bool] = self._loop.create_future()
|
|
|
|
def _done(ok: bool, _nodes: Iterable[Any]) -> None:
|
|
self._loop.call_soon_threadsafe(fut.set_result, bool(ok))
|
|
|
|
self._dht.putSigned(key, val, _done, permanent)
|
|
return await fut
|
|
|
|
async def putEncrypted(
|
|
self,
|
|
key: InfoHash,
|
|
to: InfoHash | PublicKey | PkId,
|
|
val: Value,
|
|
*,
|
|
permanent: bool = False,
|
|
) -> bool:
|
|
fut: asyncio.Future[bool] = self._loop.create_future()
|
|
|
|
def _done(ok: bool, _nodes: Iterable[Any]) -> None:
|
|
self._loop.call_soon_threadsafe(fut.set_result, bool(ok))
|
|
|
|
self._dht.putEncrypted(key, val, to, _done, permanent)
|
|
return await fut
|
|
|
|
def listen(self, key: InfoHash) -> _Listener:
|
|
return _Listener(self, key)
|
|
|
|
async def cancelListen(self, listener: _Listener) -> None:
|
|
await listener.cancel()
|
|
|
|
async def shutdown(self) -> None:
|
|
fut: asyncio.Future[None] = self._loop.create_future()
|
|
self._dht.shutdown(lambda: self._loop.call_soon_threadsafe(fut.set_result, None))
|
|
await fut
|
|
|
|
async def join(self) -> None:
|
|
await asyncio.to_thread(self._dht.join)
|
|
|
|
async def __aenter__(self):
|
|
if not self.isRunning():
|
|
self.run()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
if self.isRunning():
|
|
try:
|
|
await self.shutdown()
|
|
finally:
|
|
await self.join()
|
|
return False
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
return getattr(self._dht, name)
|