koi-net 1.0.0b19__py3-none-any.whl → 1.1.0b2__py3-none-any.whl

koi_net/processor/knowledge_pipeline.py

@@ -0,0 +1,220 @@
+import logging
+from typing import Callable
+from rid_lib.core import RIDType
+from rid_lib.types import KoiNetEdge, KoiNetNode
+from rid_lib.ext import Cache
+from ..protocol.event import EventType
+from ..network.request_handler import RequestHandler
+from ..network.event_queue import NetworkEventQueue
+from ..network.graph import NetworkGraph
+from ..identity import NodeIdentity
+from .handler import (
+    KnowledgeHandler,
+    HandlerType, 
+    STOP_CHAIN,
+    StopChain
+)
+from .knowledge_object import KnowledgeObject
+
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from ..context import HandlerContext
+
+logger = logging.getLogger(__name__)
+
+
+class KnowledgePipeline:
+    handler_context: "HandlerContext"
+    cache: Cache
+    identity: NodeIdentity
+    request_handler: RequestHandler
+    event_queue: NetworkEventQueue
+    graph: NetworkGraph
+    handlers: list[KnowledgeHandler]
+    
+    def __init__(
+        self, 
+        handler_context: "HandlerContext",
+        cache: Cache, 
+        request_handler: RequestHandler,
+        event_queue: NetworkEventQueue,
+        graph: NetworkGraph,
+        default_handlers: list[KnowledgeHandler] = []
+    ):
+        self.handler_context = handler_context
+        self.cache = cache
+        self.request_handler = request_handler
+        self.event_queue = event_queue
+        self.graph = graph
+        self.handlers = default_handlers
+    
+    def add_handler(self, handler: KnowledgeHandler):
+        self.handlers.append(handler)
+            
+    def register_handler(
+        self,
+        handler_type: HandlerType,
+        rid_types: list[RIDType] | None = None,
+        event_types: list[EventType | None] | None = None
+    ):
+        """Assigns decorated function as handler for this processor."""
+        def decorator(func: Callable) -> Callable:
+            handler = KnowledgeHandler(func, handler_type, rid_types, event_types)
+            self.add_handler(handler)
+            return func
+        return decorator
+            
+    def call_handler_chain(
+        self, 
+        handler_type: HandlerType,
+        kobj: KnowledgeObject
+    ) -> KnowledgeObject | StopChain:
+        """Calls handlers of provided type, chaining their inputs and outputs together.
+        
+        The knowledge object provided when this function is called will be passed to the first handler. A handler may return one of three types: 
+        - `KnowledgeObject` - to modify the knowledge object for the next handler in the chain
+        - `None` - to keep the same knowledge object for the next handler in the chain
+        - `STOP_CHAIN` - to stop the handler chain and immediately exit the processing pipeline
+        
+        Handlers will only be called in the chain if their handler and RID type match that of the inputted knowledge object. 
+        """
+        
+        for handler in self.handlers:
+            if handler_type != handler.handler_type: 
+                continue
+            
+            if handler.rid_types and type(kobj.rid) not in handler.rid_types:
+                continue
+            
+            if handler.event_types and kobj.event_type not in handler.event_types:
+                continue
+            
+            logger.debug(f"Calling {handler_type} handler '{handler.func.__name__}'")
+            
+            resp = handler.func(
+                ctx=self.handler_context,
+                kobj=kobj.model_copy()
+            )
+            
+            # stops handler chain execution
+            if resp is STOP_CHAIN:
+                logger.debug(f"Handler chain stopped by {handler.func.__name__}")
+                return STOP_CHAIN
+            # kobj unmodified
+            elif resp is None:
+                continue
+            # kobj modified by handler
+            elif isinstance(resp, KnowledgeObject):
+                kobj = resp
+                logger.debug(f"Knowledge object modified by {handler.func.__name__}")
+            else:
+                raise ValueError(f"Handler {handler.func.__name__} returned invalid response '{resp}'")
+                    
+        return kobj
+    
+    def process(self, kobj: KnowledgeObject):
+        """Sends provided knowledge obejct through knowledge processing pipeline.
+        
+        Handler chains are called in between major events in the pipeline, indicated by their handler type. Each handler type is guaranteed to have access to certain knowledge, and may affect a subsequent action in the pipeline. The five handler types are as follows:
+        - RID - provided RID; if event type is `FORGET`, this handler decides whether to delete the knowledge from the cache by setting the normalized event type to `FORGET`, otherwise this handler decides whether to validate the manifest (and fetch it if not provided).
+        - Manifest - provided RID, manifest; decides whether to validate the bundle (and fetch it if not provided).
+        - Bundle - provided RID, manifest, contents (bundle); decides whether to write knowledge to the cache by setting the normalized event type to `NEW` or `UPDATE`.
+        - Network - provided RID, manifest, contents (bundle); decides which nodes (if any) to broadcast an event about this knowledge to. (Note, if event type is `FORGET`, the manifest and contents will be retrieved from the local cache, and indicate the last state of the knowledge before it was deleted.)
+        - Final - provided RID, manifests, contents (bundle); final action taken after network broadcast.
+        
+        The pipeline may be stopped by any point by a single handler returning the `STOP_CHAIN` sentinel. In that case, the process will exit immediately. Further handlers of that type and later handler chains will not be called.
+        """
+        
+        logger.debug(f"Handling {kobj!r}")
+        kobj = self.call_handler_chain(HandlerType.RID, kobj)
+        if kobj is STOP_CHAIN: return
+        
+        if kobj.event_type == EventType.FORGET:
+            bundle = self.cache.read(kobj.rid)
+            if not bundle: 
+                logger.debug("Local bundle not found")
+                return
+            
+            # the bundle (to be deleted) attached to kobj for downstream analysis
+            logger.debug("Adding local bundle (to be deleted) to knowledge object")
+            kobj.manifest = bundle.manifest
+            kobj.contents = bundle.contents
+            
+        else:
+            # attempt to retrieve manifest
+            if not kobj.manifest:
+                logger.debug("Manifest not found")
+                if not kobj.source:
+                    return
+            
+                logger.debug("Attempting to fetch remote manifest from source")
+                payload = self.request_handler.fetch_manifests(
+                    node=kobj.source,
+                    rids=[kobj.rid]
+                )
+                    
+                if not payload.manifests:
+                    logger.debug("Failed to find manifest")
+                    return
+                
+                kobj.manifest = payload.manifests[0]
+                
+            kobj = self.call_handler_chain(HandlerType.Manifest, kobj)
+            if kobj is STOP_CHAIN: return
+            
+            # attempt to retrieve bundle
+            if not kobj.bundle:
+                logger.debug("Bundle not found")
+                if kobj.source is None:
+                    return
+                
+                logger.debug("Attempting to fetch remote bundle from source")
+                payload = self.request_handler.fetch_bundles(
+                    node=kobj.source,
+                    rids=[kobj.rid]
+                )
+                
+                if not payload.bundles:
+                    logger.debug("Failed to find bundle")
+                    return
+                
+                bundle = payload.bundles[0]                    
+                
+                if kobj.manifest != bundle.manifest:
+                    logger.warning("Retrieved bundle contains a different manifest")
+                
+                kobj.manifest = bundle.manifest
+                kobj.contents = bundle.contents                
+                
+        kobj = self.call_handler_chain(HandlerType.Bundle, kobj)
+        if kobj is STOP_CHAIN: return
+            
+        if kobj.normalized_event_type in (EventType.UPDATE, EventType.NEW):
+            logger.info(f"Writing to cache: {kobj!r}")
+            self.cache.write(kobj.bundle)
+            
+        elif kobj.normalized_event_type == EventType.FORGET:
+            logger.info(f"Deleting from cache: {kobj!r}")
+            self.cache.delete(kobj.rid)
+            
+        else:
+            logger.debug("Normalized event type was never set, no cache or network operations will occur")
+            return
+        
+        if type(kobj.rid) in (KoiNetNode, KoiNetEdge):
+            logger.debug("Change to node or edge, regenerating network graph")
+            self.graph.generate()
+        
+        kobj = self.call_handler_chain(HandlerType.Network, kobj)
+        if kobj is STOP_CHAIN: return
+        
+        if kobj.network_targets:
+            logger.debug(f"Broadcasting event to {len(kobj.network_targets)} network target(s)")
+        else:
+            logger.debug("No network targets set")
+        
+        for node in kobj.network_targets:
+            self.event_queue.push_event_to(kobj.normalized_event, node)
+            self.event_queue.flush_webhook_queue(node)
+        
+        kobj = self.call_handler_chain(HandlerType.Final, kobj)

koi_net/protocol/api_models.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel
 from rid_lib import RID, RIDType
 from rid_lib.ext import Bundle, Manifest
 from .event import Event
+from .errors import ErrorTypes
 
 
 # REQUEST MODELS
@@ -41,7 +42,12 @@ class EventsPayload(BaseModel):
     events: list[Event]
     
 
+# ERROR MODELS
+
+class ErrorResponse(BaseModel):
+    error: ErrorTypes
+
 # TYPES
 
 type RequestModels = EventsPayload | PollEvents | FetchRids | FetchManifests | FetchBundles
-type ResponseModels = RidsPayload | ManifestsPayload | BundlesPayload | EventsPayload
+type ResponseModels = RidsPayload | ManifestsPayload | BundlesPayload | EventsPayload | ErrorResponse

koi_net/protocol/edge.py

@@ -1,7 +1,8 @@
 from enum import StrEnum
 from pydantic import BaseModel
 from rid_lib import RIDType
-from rid_lib.types import KoiNetNode
+from rid_lib.ext.bundle import Bundle
+from rid_lib.types import KoiNetEdge, KoiNetNode
 
 
 class EdgeStatus(StrEnum):
@@ -18,3 +19,24 @@ class EdgeProfile(BaseModel):
     edge_type: EdgeType
     status: EdgeStatus
     rid_types: list[RIDType]
+
+
+def generate_edge_bundle(
+    source: KoiNetNode,
+    target: KoiNetNode,
+    rid_types: list[RIDType],
+    edge_type: EdgeType
+) -> Bundle:
+    edge_rid = KoiNetEdge.generate(source, target)
+    edge_profile = EdgeProfile(
+        source=source,
+        target=target,
+        rid_types=rid_types,
+        edge_type=edge_type,
+        status=EdgeStatus.PROPOSED
+    )
+    edge_bundle = Bundle.generate(
+        edge_rid,
+        edge_profile.model_dump()
+    )
+    return edge_bundle

koi_net/protocol/envelope.py

@@ -0,0 +1,54 @@
+import logging
+from typing import Generic, TypeVar
+from pydantic import BaseModel
+from rid_lib.types import KoiNetNode
+
+from .secure import PrivateKey, PublicKey
+from .api_models import RequestModels, ResponseModels
+
+
+logger = logging.getLogger(__name__)
+
+
+T = TypeVar("T", bound=RequestModels | ResponseModels)
+
+class SignedEnvelope(BaseModel, Generic[T]):
+    payload: T
+    source_node: KoiNetNode
+    target_node: KoiNetNode
+    signature: str
+    
+    def verify_with(self, pub_key: PublicKey):        
+        # IMPORTANT: calling `model_dump()` loses all typing! when converting between SignedEnvelope and UnsignedEnvelope, use the Pydantic classes, not the dictionary form
+        unsigned_envelope = UnsignedEnvelope[T](
+            payload=self.payload,
+            source_node=self.source_node,
+            target_node=self.target_node 
+        )
+        
+        logger.debug(f"Verifying envelope: {unsigned_envelope.model_dump_json()}")
+        
+        pub_key.verify(
+            self.signature,
+            unsigned_envelope.model_dump_json().encode()
+        )
+
+class UnsignedEnvelope(BaseModel, Generic[T]):
+    payload: T
+    source_node: KoiNetNode
+    target_node: KoiNetNode
+    
+    def sign_with(self, priv_key: PrivateKey) -> SignedEnvelope[T]:
+        logger.debug(f"Signing envelope: {self.model_dump_json()}")
+        logger.debug(f"Type: [{type(self.payload)}]")
+        
+        signature = priv_key.sign(
+            self.model_dump_json().encode()
+        )
+        
+        return SignedEnvelope(
+            payload=self.payload,
+            source_node=self.source_node,
+            target_node=self.target_node,
+            signature=signature
+        )

koi_net/protocol/errors.py

@@ -0,0 +1,23 @@
+from enum import StrEnum
+
+
+class ErrorTypes(StrEnum):
+    UnknownNode = "unknown_node"
+    InvalidKey = "invalid_key"
+    InvalidSignature = "invalid_signature"
+    InvalidTarget = "invalid_target"
+
+class ProtocolError(Exception):
+    error_type: ErrorTypes
+    
+class UnknownNodeError(ProtocolError):
+    error_type = ErrorTypes.UnknownNode
+    
+class InvalidKeyError(ProtocolError):
+    error_type = ErrorTypes.InvalidKey
+    
+class InvalidSignatureError(ProtocolError):
+    error_type = ErrorTypes.InvalidSignature
+
+class InvalidTargetError(ProtocolError):
+    error_type = ErrorTypes.InvalidTarget

koi_net/protocol/node.py

@@ -14,4 +14,5 @@ class NodeProvides(BaseModel):
 class NodeProfile(BaseModel):
     base_url: str | None = None
     node_type: NodeType
-    provides: NodeProvides = NodeProvides()
+    provides: NodeProvides = NodeProvides()
+    public_key: str | None = None

koi_net/protocol/secure.py

@@ -0,0 +1,106 @@
+import logging
+from base64 import urlsafe_b64decode, urlsafe_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
+
+logger = logging.getLogger(__name__)
+
+
+class PrivateKey:
+    priv_key: ec.EllipticCurvePrivateKey
+    
+    def __init__(self, priv_key):
+        self.priv_key = priv_key
+        
+    @classmethod
+    def generate(cls):
+        return cls(priv_key=ec.generate_private_key(ec.SECP192R1()))
+
+    def public_key(self) -> "PublicKey":
+        return PublicKey(self.priv_key.public_key())
+    
+    @classmethod
+    def from_pem(cls, priv_key_pem: str, password: str):
+        return cls(
+            priv_key=serialization.load_pem_private_key(
+                data=priv_key_pem.encode(),
+                password=password.encode()
+            )
+        )
+
+    def to_pem(self, password: str) -> str:
+        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:
+        hashed_message = sha256_hash(message.decode())
+        
+        signature = urlsafe_b64encode(
+            self.priv_key.sign(
+                data=message,
+                signature_algorithm=ec.ECDSA(hashes.SHA256())
+            )
+        ).decode()
+        
+        logger.debug(f"Signing message with [{self.public_key().to_der()}]")
+        logger.debug(f"hash: {hashed_message}")
+        logger.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):
+        return cls(
+            pub_key=serialization.load_pem_public_key(
+                data=pub_key_pem.encode()
+            )
+        )
+        
+    def to_pem(self) -> str:
+        return self.pub_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        ).decode()
+        
+    @classmethod
+    def from_der(cls, pub_key_der: str):        
+        return cls(
+            pub_key=serialization.load_der_public_key(
+                data=urlsafe_b64decode(pub_key_der)
+            )
+        )
+    
+    def to_der(self) -> str:
+        return urlsafe_b64encode(
+            self.pub_key.public_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PublicFormat.SubjectPublicKeyInfo
+            )
+        ).decode()
+        
+    def verify(self, signature: str, message: bytes) -> bool:
+        hashed_message = sha256_hash(message.decode())
+
+        logger.debug(f"Verifying message with [{self.to_der()}]")
+        logger.debug(f"hash: {hashed_message}")
+        logger.debug(f"signature: {signature}")
+        
+        # NOTE: throws cryptography.exceptions.InvalidSignature on failure
+        
+        self.pub_key.verify(
+            signature=urlsafe_b64decode(signature),
+            data=message,
+            signature_algorithm=ec.ECDSA(hashes.SHA256())
+        )

koi_net/secure.py

@@ -0,0 +1,117 @@
+import logging
+from functools import wraps
+
+import cryptography.exceptions
+from rid_lib.ext import Bundle
+from rid_lib.ext.utils import sha256_hash
+from .identity import NodeIdentity
+from .protocol.envelope import UnsignedEnvelope, SignedEnvelope
+from .protocol.secure import PublicKey
+from .protocol.api_models import EventsPayload
+from .protocol.event import EventType
+from .protocol.node import NodeProfile
+from .protocol.secure import PrivateKey
+from .protocol.errors import (
+    UnknownNodeError,
+    InvalidKeyError,
+    InvalidSignatureError,
+    InvalidTargetError
+)
+from .effector import Effector
+from .config import NodeConfig
+
+logger = logging.getLogger(__name__)
+
+
+class Secure:
+    identity: NodeIdentity
+    effector: Effector
+    config: NodeConfig
+    priv_key: PrivateKey
+    
+    def __init__(
+        self, 
+        identity: NodeIdentity, 
+        effector: Effector,
+        config: NodeConfig
+    ):
+        self.identity = identity
+        self.effector = effector
+        self.config = config
+
+        self.priv_key = self._load_priv_key()
+        
+    def _load_priv_key(self) -> PrivateKey:
+        with open(self.config.koi_net.private_key_pem_path, "r") as f:
+            priv_key_pem = f.read()
+        
+        return 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:
+        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, target) -> SignedEnvelope:
+        return UnsignedEnvelope(
+            payload=payload,
+            source_node=self.identity.rid,
+            target_node=target
+        ).sign_with(self.priv_key)
+        
+    def validate_envelope(self, envelope: SignedEnvelope):
+        node_bundle = (
+            self.effector.deref(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 as err:
+            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")
+        
+    def envelope_handler(self, func):
+        @wraps(func)
+        async def wrapper(req: SignedEnvelope, *args, **kwargs) -> SignedEnvelope | None:
+            logger.info("Validating envelope")
+            
+            self.validate_envelope(req)            
+            logger.info("Calling endpoint handler")
+            
+            result = await func(req, *args, **kwargs)            
+            
+            if result is not None:
+                logger.info("Creating response envelope")
+                return self.create_envelope(
+                    payload=result,
+                    target=req.source_node
+                )
+        return wrapper
+