koi-net 1.0.0b19__py3-none-any.whl → 1.1.0__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

@@ -1,47 +1,62 @@
 """Pydantic models for request and response/payload objects in the KOI-net API."""
 
-from pydantic import BaseModel
+from typing import Literal
+from pydantic import BaseModel, Field
 from rid_lib import RID, RIDType
 from rid_lib.ext import Bundle, Manifest
 from .event import Event
+from .errors import ErrorType
 
 
 # REQUEST MODELS
 
 class PollEvents(BaseModel):
-    rid: RID
+    type: Literal["poll_events"] = Field("poll_events")
     limit: int = 0
     
 class FetchRids(BaseModel):
+    type: Literal["fetch_rids"] = Field("fetch_rids")
     rid_types: list[RIDType] = []
     
 class FetchManifests(BaseModel):
+    type: Literal["fetch_manifests"] = Field("fetch_manifests")
     rid_types: list[RIDType] = []
     rids: list[RID] = []
     
 class FetchBundles(BaseModel):
+    type: Literal["fetch_bundles"] = Field("fetch_bundles")
     rids: list[RID]
     
 
 # RESPONSE/PAYLOAD MODELS
 
 class RidsPayload(BaseModel):
+    type: Literal["rids_payload"] = Field("rids_payload")
     rids: list[RID]
 
 class ManifestsPayload(BaseModel):
+    type: Literal["manifests_payload"] = Field("manifests_payload")
     manifests: list[Manifest]
     not_found: list[RID] = []
     
 class BundlesPayload(BaseModel):
+    type: Literal["bundles_payload"] = Field("bundles_payload")
     bundles: list[Bundle]
     not_found: list[RID] = []
     deferred: list[RID] = []
     
 class EventsPayload(BaseModel):
+    type: Literal["events_payload"] = Field("events_payload")
     events: list[Event]
     
 
+# ERROR MODELS
+
+class ErrorResponse(BaseModel):
+    type: Literal["error_response"] = Field("error_response")
+    error: ErrorType
+
 # 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,9 @@
 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.ext.utils import sha256_hash
+from rid_lib.types import KoiNetEdge, KoiNetNode
 
 
 class EdgeStatus(StrEnum):
@@ -18,3 +20,26 @@ 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(sha256_hash(
+        str(source) + str(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,58 @@
+import logging
+from typing import Generic, TypeVar
+from pydantic import BaseModel, ConfigDict
+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]):
+    model_config = ConfigDict(exclude_none=True)
+    
+    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(exclude_none=True)}")
+                
+        pub_key.verify(
+            self.signature,
+            unsigned_envelope.model_dump_json(exclude_none=True).encode()
+        )
+
+class UnsignedEnvelope(BaseModel, Generic[T]):
+    model_config = ConfigDict(exclude_none=True)
+    
+    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(exclude_none=True)}")
+        logger.debug(f"Type: [{type(self.payload)}]")
+        
+        signature = priv_key.sign(
+            self.model_dump_json(exclude_none=True).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 ErrorType(StrEnum):
+    UnknownNode = "unknown_node"
+    InvalidKey = "invalid_key"
+    InvalidSignature = "invalid_signature"
+    InvalidTarget = "invalid_target"
+
+class ProtocolError(Exception):
+    error_type: ErrorType
+    
+class UnknownNodeError(ProtocolError):
+    error_type = ErrorType.UnknownNode
+    
+class InvalidKeyError(ProtocolError):
+    error_type = ErrorType.InvalidKey
+    
+class InvalidSignatureError(ProtocolError):
+    error_type = ErrorType.InvalidSignature
+
+class InvalidTargetError(ProtocolError):
+    error_type = ErrorType.InvalidTarget

koi_net/protocol/event.py

@@ -15,9 +15,6 @@ class Event(BaseModel):
     manifest: Manifest | None = None
     contents: dict | None = None
     
-    class Config:
-        exclude_none = True
-    
     def __repr__(self):
         return f"<Event '{self.rid}' event type: '{self.event_type}'>"
     

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,160 @@
+import logging
+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
+)
+
+logger = logging.getLogger(__name__)
+
+
+def der_to_raw_signature(der_signature: bytes, curve=ec.SECP256R1()) -> bytes:
+    """Convert 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:
+    """Convert 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):
+        return cls(priv_key=ec.generate_private_key(ec.SECP256R1()))
+
+    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())
+        
+        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()
+        
+        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=b64decode(pub_key_der)
+            )
+        )
+    
+    def to_der(self) -> str:
+        return 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())
+        
+        # print(message.hex())
+        # print()
+        # print(hashed_message)
+        # print()
+        # print(message.decode())
+
+        # logger.debug(f"Verifying message with [{self.to_der()}]")
+        # logger.debug(f"hash: {hashed_message}")
+        # logger.debug(f"signature: {signature}")
+        
+        raw_signature_bytes = b64decode(signature)
+        der_signature_bytes = raw_to_der_signature(raw_signature_bytes)
+        
+        # NOTE: throws cryptography.exceptions.InvalidSignature on failure
+        
+        self.pub_key.verify(
+            signature=der_signature_bytes,
+            data=message,
+            signature_algorithm=ec.ECDSA(hashes.SHA256())
+        )