koi-net 1.2.4__py3-none-any.whl

koi_net/protocol/secure.py

@@ -0,0 +1,167 @@
+from rid_lib.types import KoiNetNode
+import structlog
+from base64 import b64decode, b64encode
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives import serialization
+from rid_lib.ext.utils import sha256_hash
+from cryptography.hazmat.primitives.asymmetric.utils import (
+    decode_dss_signature, 
+    encode_dss_signature
+)
+
+log = structlog.stdlib.get_logger()
+
+
+def der_to_raw_signature(der_signature: bytes, curve=ec.SECP256R1()) -> bytes:
+    """Converts a DER-encoded signature to raw r||s format."""
+    
+    # Decode the DER signature to get r and s
+    r, s = decode_dss_signature(der_signature)
+    
+    # Determine byte length based on curve bit size
+    byte_length = (curve.key_size + 7) // 8
+    
+    # Convert r and s to big-endian byte arrays of fixed length
+    r_bytes = r.to_bytes(byte_length, byteorder='big')
+    s_bytes = s.to_bytes(byte_length, byteorder='big')
+    
+    # Concatenate r and s
+    return r_bytes + s_bytes
+
+
+def raw_to_der_signature(raw_signature: bytes, curve=ec.SECP256R1()) -> bytes:
+    """Converts a raw r||s signature to DER format."""
+    
+    # Determine byte length based on curve bit size
+    byte_length = (curve.key_size + 7) // 8
+    
+    # Split the raw signature into r and s components
+    if len(raw_signature) != 2 * byte_length:
+        raise ValueError(f"Raw signature must be {2 * byte_length} bytes for {curve.name}")
+    
+    r_bytes = raw_signature[:byte_length]
+    s_bytes = raw_signature[byte_length:]
+    
+    # Convert bytes to integers
+    r = int.from_bytes(r_bytes, byteorder='big')
+    s = int.from_bytes(s_bytes, byteorder='big')
+    
+    # Encode as DER
+    return encode_dss_signature(r, s)
+
+
+class PrivateKey:
+    priv_key: ec.EllipticCurvePrivateKey
+    
+    def __init__(self, priv_key):
+        self.priv_key = priv_key
+    
+    @classmethod
+    def generate(cls):
+        """Generates a new `Private Key`."""
+        return cls(priv_key=ec.generate_private_key(ec.SECP256R1()))
+
+    def public_key(self) -> "PublicKey":
+        """Returns instance of `PublicKey` dervied from this private key."""
+        return PublicKey(self.priv_key.public_key())
+    
+    @classmethod
+    def from_pem(cls, priv_key_pem: str, password: str):
+        """Loads `PrivateKey` from encrypted PEM string."""
+        return cls(
+            priv_key=serialization.load_pem_private_key(
+                data=priv_key_pem.encode(),
+                password=password.encode()
+            )
+        )
+
+    def to_pem(self, password: str) -> str:
+        """Saves `PrivateKey` to encrypted PEM string."""
+        return self.priv_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.BestAvailableEncryption(password.encode())
+        ).decode()
+        
+    def sign(self, message: bytes) -> str:
+        """Returns base64 encoded raw signature bytes of the form r||s."""
+        hashed_message = sha256_hash(message.decode())
+        
+        der_signature_bytes = self.priv_key.sign(
+            data=message,
+            signature_algorithm=ec.ECDSA(hashes.SHA256())
+        )
+        
+        raw_signature_bytes = der_to_raw_signature(der_signature_bytes)
+        
+        signature = b64encode(raw_signature_bytes).decode()
+        
+        log.debug(f"Signing message with [{self.public_key().to_der()}]")
+        log.debug(f"hash: {hashed_message}")
+        log.debug(f"signature: {signature}")
+        
+        return signature
+
+
+class PublicKey:
+    pub_key: ec.EllipticCurvePublicKey
+    
+    def __init__(self, pub_key):
+        self.pub_key = pub_key
+    
+    @classmethod
+    def from_pem(cls, pub_key_pem: str):
+        """Loads `PublicKey` from PEM string."""
+        return cls(
+            pub_key=serialization.load_pem_public_key(
+                data=pub_key_pem.encode()
+            )
+        )
+        
+    def to_pem(self) -> str:
+        """Saves `PublicKey` to PEM string."""
+        return self.pub_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        ).decode()
+        
+    @classmethod
+    def from_der(cls, pub_key_der: str):
+        """Loads `PublicKey` from base64 encoded DER string."""
+        return cls(
+            pub_key=serialization.load_der_public_key(
+                data=b64decode(pub_key_der)
+            )
+        )
+    
+    def to_der(self) -> str:
+        """Saves `PublicKey` to base64 encoded DER string."""
+        return b64encode(
+            self.pub_key.public_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PublicFormat.SubjectPublicKeyInfo
+            )
+        ).decode()
+        
+    def to_node_rid(self, name) -> KoiNetNode:
+        """Returns an orn:koi-net.node RID from hashed DER string."""
+        return KoiNetNode(
+            name=name,
+            hash=sha256_hash(self.to_der())
+        )
+        
+    def verify(self, signature: str, message: bytes):
+        """Verifies a signature for a message.
+        
+        Raises `cryptography.exceptions.InvalidSignature` on failure.
+        """
+        
+        raw_signature_bytes = b64decode(signature)
+        der_signature_bytes = raw_to_der_signature(raw_signature_bytes)
+        
+        self.pub_key.verify(
+            signature=der_signature_bytes,
+            data=message,
+            signature_algorithm=ec.ECDSA(hashes.SHA256())
+        )

koi_net/secure_manager.py

@@ -0,0 +1,115 @@
+import structlog
+import cryptography.exceptions
+from rid_lib.ext import Bundle, Cache
+from rid_lib.ext.utils import sha256_hash
+from rid_lib.types import KoiNetNode
+from .identity import NodeIdentity
+from .protocol.envelope import UnsignedEnvelope, SignedEnvelope
+from .protocol.secure import PublicKey
+from .protocol.api_models import ApiModels, EventsPayload
+from .protocol.event import EventType
+from .protocol.node import NodeProfile
+from .protocol.secure import PrivateKey
+from .exceptions import (
+    UnknownNodeError,
+    InvalidKeyError,
+    InvalidSignatureError,
+    InvalidTargetError
+)
+from .config.core import NodeConfig
+
+log = structlog.stdlib.get_logger()
+
+
+class SecureManager:
+    """Subsystem handling secure protocol logic."""
+    identity: NodeIdentity
+    cache: Cache
+    config: NodeConfig
+    priv_key: PrivateKey
+    
+    def __init__(
+        self, 
+        identity: NodeIdentity, 
+        cache: Cache,
+        config: NodeConfig
+    ):
+        self.identity = identity
+        self.cache = cache
+        self.config = config
+        
+    def start(self):
+        self.load_priv_key()
+        
+    def load_priv_key(self) -> PrivateKey:
+        """Loads private key from PEM file path in config."""
+        
+        # TODO: handle missing private key
+        with open(self.config.koi_net.private_key_pem_path, "r") as f:
+            priv_key_pem = f.read()
+        
+        self.priv_key = PrivateKey.from_pem(
+            priv_key_pem=priv_key_pem,
+            password=self.config.env.priv_key_password
+        )
+        
+    def handle_unknown_node(self, envelope: SignedEnvelope) -> Bundle | None:
+        """Attempts to find node profile in proided envelope.
+        
+        If an unknown node sends an envelope, it may still be able to be
+        validated if that envelope contains their node profile. This is
+        essential for allowing unknown nodes to handshake and introduce
+        themselves. Only an `EventsPayload` contain a `NEW` event for a 
+        node profile for the source node is permissible.
+        """
+        if type(envelope.payload) != EventsPayload:
+            return None
+            
+        for event in envelope.payload.events:
+            # must be NEW event for bundle of source node's profile
+            if event.rid != envelope.source_node:
+                continue
+            if event.event_type != EventType.NEW:
+                continue            
+            
+            return event.bundle
+        return None
+        
+    def create_envelope(
+        self, payload: ApiModels, target: KoiNetNode
+    ) -> SignedEnvelope:
+        """Returns signed envelope to target from provided payload."""
+        return UnsignedEnvelope(
+            payload=payload,
+            source_node=self.identity.rid,
+            target_node=target
+        ).sign_with(self.priv_key)
+        
+    def validate_envelope(self, envelope: SignedEnvelope):
+        """Validates signed envelope from another node."""
+        
+        node_bundle = (
+            self.cache.read(envelope.source_node) or
+            self.handle_unknown_node(envelope)
+        )
+        
+        if not node_bundle:
+            raise UnknownNodeError(f"Couldn't resolve {envelope.source_node}")
+        
+        node_profile = node_bundle.validate_contents(NodeProfile)
+        
+        # check that public key matches source node RID
+        if envelope.source_node.hash != sha256_hash(node_profile.public_key):
+            raise InvalidKeyError("Invalid public key on new node!")
+        
+        # check envelope signed by validated public key
+        pub_key = PublicKey.from_der(node_profile.public_key)
+        try:
+            envelope.verify_with(pub_key)
+        except cryptography.exceptions.InvalidSignature:
+            raise InvalidSignatureError(f"Signature {envelope.signature} is invalid.")
+        
+        # check that this node is the target of the envelope
+        if envelope.target_node != self.identity.rid:
+            raise InvalidTargetError(f"Envelope target {envelope.target_node!r} is not me")
+

koi_net/workers/__init__.py

@@ -0,0 +1,2 @@
+from .event_worker import EventProcessingWorker
+from .kobj_worker import KnowledgeProcessingWorker

koi_net/workers/base.py

@@ -0,0 +1,26 @@
+import threading
+
+from ..build import comp_order
+
+
+class End:
+    """Class for STOP_WORKER sentinel pushed to worker queues."""
+    pass
+
+STOP_WORKER = End()
+
+@comp_order.worker
+class ThreadWorker:
+    """Base class for thread workers."""
+    
+    thread: threading.Thread
+    
+    def __init__(self):
+        self.thread = threading.Thread(target=self.run)
+        
+    def start(self):
+        self.thread.start()
+        
+    def run(self):
+        """Processing loop for thread."""
+        pass

koi_net/workers/event_worker.py

@@ -0,0 +1,111 @@
+import queue
+import traceback
+import time
+import structlog
+
+from rid_lib.ext import Cache
+from rid_lib.types import KoiNetNode
+
+from ..config.core import NodeConfig
+from ..network.event_queue import EventQueue
+from ..network.request_handler import RequestHandler
+from ..network.event_buffer import EventBuffer
+from ..protocol.node import NodeProfile, NodeType
+from ..exceptions import RequestError
+from .base import ThreadWorker, STOP_WORKER
+
+log = structlog.stdlib.get_logger()
+
+
+class EventProcessingWorker(ThreadWorker):
+    """Thread worker that processes the `event_queue`."""
+    
+    def __init__(
+        self,
+        config: NodeConfig,
+        cache: Cache,
+        event_queue: EventQueue,
+        request_handler: RequestHandler,
+        poll_event_buf: EventBuffer,
+        broadcast_event_buf: EventBuffer
+    ):
+        self.event_queue = event_queue
+        self.request_handler = request_handler
+        
+        self.config = config
+        self.cache = cache
+        self.poll_event_buf = poll_event_buf
+        self.broadcast_event_buf = broadcast_event_buf
+        
+        super().__init__()
+        
+    def flush_and_broadcast(self, target: KoiNetNode, force_flush: bool = False):
+        """Broadcasts all events to target in event buffer."""
+        
+        # TODO: deal with automated retries when unreachable node's buffer is full
+        try:
+            with self.broadcast_event_buf.safe_flush(target, force_flush) as events:
+                self.request_handler.broadcast_events(target, events=events)
+        except RequestError:
+            log.warning("Failed to reach target, event buffer reset")
+            pass
+        
+    def stop(self):
+        self.event_queue.q.put(STOP_WORKER)
+    
+    def run(self):
+        while True:
+            try:
+                item = self.event_queue.q.get(
+                    timeout=self.config.koi_net.event_worker.queue_timeout)
+                
+                try:
+                    if item is STOP_WORKER:
+                        log.info(f"Received 'STOP_WORKER' signal, flushing all buffers...")
+                        for target in list(self.broadcast_event_buf.buffers.keys()):
+                            self.flush_and_broadcast(target, force_flush=True)
+                        return
+                    
+                    log.info(f"Dequeued {item.event!r} -> {item.target!r}")
+                    
+                    # determines which buffer to push event to based on target node type
+                    node_bundle = self.cache.read(item.target)
+                    if node_bundle:
+                        node_profile = node_bundle.validate_contents(NodeProfile)
+                        
+                        if node_profile.node_type == NodeType.FULL:
+                            self.broadcast_event_buf.push(item.target, item.event)
+                            
+                        elif node_profile.node_type == NodeType.PARTIAL:
+                            self.poll_event_buf.push(item.target, item.event)
+                            continue
+                        
+                    elif item.target == self.config.koi_net.first_contact.rid:
+                        self.broadcast_event_buf.push(item.target, item.event)
+                        
+                    else:
+                        log.warning(f"Couldn't handle event {item.event!r} in queue, node {item.target!r} unknown to me")
+                        continue
+                    
+                    buf_len = self.broadcast_event_buf.buf_len(item.target)
+                    if buf_len > self.config.koi_net.event_worker.max_buf_len:
+                        self.flush_and_broadcast(target)
+
+                finally:
+                    self.event_queue.q.task_done()
+
+            except queue.Empty:
+                # On timeout, check all buffers for max wait time
+                for target in list(self.broadcast_event_buf.buffers):
+                    start_time = self.broadcast_event_buf.start_time.get(target)
+                    
+                    if (start_time is None) or (self.broadcast_event_buf.buf_len(target) == 0):
+                        continue
+                    
+                    now = time.time()
+                    if (now - start_time) >= self.config.koi_net.event_worker.max_wait_time: 
+                        self.flush_and_broadcast(target)
+                        
+            except Exception:
+                traceback.print_exc()
+                continue

koi_net/workers/kobj_worker.py

@@ -0,0 +1,51 @@
+import queue
+import traceback
+import structlog
+
+from ..config.core import NodeConfig
+from ..processor.pipeline import KnowledgePipeline
+from ..processor.kobj_queue import KobjQueue
+from .base import ThreadWorker, STOP_WORKER
+
+log = structlog.stdlib.get_logger()
+
+
+class KnowledgeProcessingWorker(ThreadWorker):
+    """Thread worker that processes the `kobj_queue`."""
+    
+    def __init__(
+        self,
+        config: NodeConfig,
+        kobj_queue: KobjQueue,
+        pipeline: KnowledgePipeline
+    ):
+        self.config = config
+        self.kobj_queue = kobj_queue
+        self.pipeline = pipeline
+
+        super().__init__()
+        
+    def stop(self):
+        self.kobj_queue.q.put(STOP_WORKER)
+        
+    def run(self):
+        while True:
+            try:
+                item = self.kobj_queue.q.get(timeout=self.config.koi_net.kobj_worker.queue_timeout)
+                try:
+                    if item is STOP_WORKER:
+                        log.info("Received 'STOP_WORKER' signal, shutting down...")
+                        return
+                    
+                    log.info(f"Dequeued {item!r}")
+                    
+                    self.pipeline.process(item)
+                finally:
+                    self.kobj_queue.q.task_done()
+                    
+            except queue.Empty:
+                pass
+            
+            except Exception:
+                traceback.print_exc()
+                continue