add typing to security

libp2p/network/connection/raw_connection_interface.py

@@ -1,4 +1,4 @@
-from abc import ABC
+from abc import ABC, abstractmethod
 
 # pylint: disable=too-few-public-methods
 
@@ -7,3 +7,11 @@ class IRawConnection(ABC):
     """
     A Raw Connection provides a Reader and a Writer
     """
+
+    @abstractmethod
+    async def write(self, data: bytes) -> None:
+        pass
+
+    @abstractmethod
+    async def read(self) -> bytes:
+        pass

libp2p/security/insecure_security.py

@@ -1,12 +1,22 @@
 from libp2p.security.secure_transport_interface import ISecureTransport
 from libp2p.security.secure_conn_interface import ISecureConn
 
+from typing import TYPE_CHECKING, Dict, Any, cast
+
+if TYPE_CHECKING:
+    from .secure_conn_interface import ISecureConn
+    from libp2p.network.connection.raw_connection_interface import IRawConnection
+    from libp2p.peer.id import ID
+    from .typing import TSecurityDetails
+
 
 class InsecureTransport(ISecureTransport):
-    def __init__(self, transport_id):
+    transport_id: int
+
+    def __init__(self, transport_id: int) -> None:
         self.transport_id = transport_id
 
-    async def secure_inbound(self, conn):
+    async def secure_inbound(self, conn: "IRawConnection") -> ISecureConn:
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are not the initiator)
@@ -15,7 +25,9 @@ class InsecureTransport(ISecureTransport):
         insecure_conn = InsecureConn(conn, self.transport_id)
         return insecure_conn
 
-    async def secure_outbound(self, conn, peer_id):
+    async def secure_outbound(
+        self, conn: "IRawConnection", peer_id: "ID"
+    ) -> ISecureConn:
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are the initiator)
@@ -26,18 +38,21 @@ class InsecureTransport(ISecureTransport):
 
 
 class InsecureConn(ISecureConn):
-    def __init__(self, conn, conn_id):
+    conn: "IRawConnection"
+    details: "TSecurityDetails"
+
+    def __init__(self, conn: "IRawConnection", conn_id: int) -> None:
         self.conn = conn
-        self.details = {}
+        self.details = cast("TSecurityDetails", {})
         self.details["id"] = conn_id
 
-    def get_conn(self):
+    def get_conn(self) -> "ISecureConn":
         """
         :return: connection object that has been made secure
         """
-        return self.conn
+        return cast("ISecureConn", self.conn)
 
-    def get_security_details(self):
+    def get_security_details(self) -> "TSecurityDetails":
         """
         :return: map containing details about the connections security
         """

libp2p/security/secure_conn_interface.py

@@ -1,5 +1,10 @@
 from abc import ABC, abstractmethod
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .typing import TSecurityDetails
+
 # pylint: disable=W0105
 
 """
@@ -12,13 +17,13 @@ Relevant go repo: https://github.com/libp2p/go-conn-security/blob/master/interfa
 
 class ISecureConn(ABC):
     @abstractmethod
-    def get_conn(self):
+    def get_conn(self) -> "ISecureConn":
         """
         :return: connection object that has been made secure
         """
 
     @abstractmethod
-    def get_security_details(self):
+    def get_security_details(self) -> "TSecurityDetails":
         """
         :return: map containing details about the connections security
         """

libp2p/security/secure_transport_interface.py

@@ -1,5 +1,12 @@
 from abc import ABC, abstractmethod
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .secure_conn_interface import ISecureConn
+    from libp2p.network.connection.raw_connection_interface import IRawConnection
+    from libp2p.peer.id import ID
+
 # pylint: disable=W0105
 
 """
@@ -12,7 +19,7 @@ Relevant go repo: https://github.com/libp2p/go-conn-security/blob/master/interfa
 
 class ISecureTransport(ABC):
     @abstractmethod
-    async def secure_inbound(self, conn):
+    async def secure_inbound(self, conn: "IRawConnection") -> "ISecureConn":
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are not the initiator)
@@ -20,7 +27,9 @@ class ISecureTransport(ABC):
         """
 
     @abstractmethod
-    async def secure_outbound(self, conn, peer_id):
+    async def secure_outbound(
+        self, conn: "IRawConnection", peer_id: "ID"
+    ) -> "ISecureConn":
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are the initiator)

libp2p/security/security_multistream.py

@@ -2,6 +2,18 @@ from abc import ABC
 from libp2p.protocol_muxer.multiselect_client import MultiselectClient
 from libp2p.protocol_muxer.multiselect import Multiselect
 
+from typing import TYPE_CHECKING, NewType, Dict
+
+if TYPE_CHECKING:
+    from libp2p.network.connection.raw_connection_interface import IRawConnection
+    from libp2p.peer.id import ID
+    from .typing import TSecurityDetails
+    from .secure_conn_interface import ISecureConn
+    from .secure_transport_interface import ISecureTransport
+
+
+TProtocol = NewType("TProtocol", str)
+
 # pylint: disable=W0105
 
 """
@@ -13,7 +25,11 @@ Relevant go repo: https://github.com/libp2p/go-conn-security/blob/master/interfa
 
 
 class SecurityMultistream(ABC):
-    def __init__(self):
+    transports: Dict[TProtocol, "ISecureTransport"]
+    multiselect: "Multiselect"
+    multiselect_client: "MultiselectClient"
+
+    def __init__(self) -> None:
         # Map protocol to secure transport
         self.transports = {}
 
@@ -23,7 +39,7 @@ class SecurityMultistream(ABC):
         # Create multiselect client
         self.multiselect_client = MultiselectClient()
 
-    def add_transport(self, protocol, transport):
+    def add_transport(self, protocol: TProtocol, transport: "ISecureTransport") -> None:
         # Associate protocol with transport
         self.transports[protocol] = transport
 
@@ -32,7 +48,7 @@ class SecurityMultistream(ABC):
         # we only care about selecting the protocol, not any handler function
         self.multiselect.add_handler(protocol, None)
 
-    async def secure_inbound(self, conn):
+    async def secure_inbound(self, conn: "IRawConnection") -> "ISecureConn":
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are not the initiator)
@@ -47,7 +63,9 @@ class SecurityMultistream(ABC):
 
         return secure_conn
 
-    async def secure_outbound(self, conn, peer_id):
+    async def secure_outbound(
+        self, conn: "IRawConnection", peer_id: "ID"
+    ) -> "ISecureConn":
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are the initiator)
@@ -62,7 +80,9 @@ class SecurityMultistream(ABC):
 
         return secure_conn
 
-    async def select_transport(self, conn, initiator):
+    async def select_transport(
+        self, conn: "IRawConnection", initiator: bool
+    ) -> "ISecureTransport":
         """
         Select a transport that both us and the node on the
         other end of conn support and agree on

libp2p/security/simple_security.py

@@ -2,12 +2,21 @@ import asyncio
 from libp2p.security.secure_transport_interface import ISecureTransport
 from libp2p.security.secure_conn_interface import ISecureConn
 
+from typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from libp2p.network.connection.raw_connection_interface import IRawConnection
+    from libp2p.peer.id import ID
+    from .typing import TSecurityDetails
+
 
 class SimpleSecurityTransport(ISecureTransport):
-    def __init__(self, key_phrase):
+    key_phrase: str
+
+    def __init__(self, key_phrase: str) -> None:
         self.key_phrase = key_phrase
 
-    async def secure_inbound(self, conn):
+    async def secure_inbound(self, conn: "IRawConnection") -> "ISecureConn":
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are not the initiator)
@@ -24,7 +33,9 @@ class SimpleSecurityTransport(ISecureTransport):
         secure_conn = SimpleSecureConn(conn, self.key_phrase)
         return secure_conn
 
-    async def secure_outbound(self, conn, peer_id):
+    async def secure_outbound(
+        self, conn: "IRawConnection", peer_id: "ID"
+    ) -> "ISecureConn":
         """
         Secure the connection, either locally or by communicating with opposing node via conn,
         for an inbound connection (i.e. we are the initiator)
@@ -47,18 +58,22 @@ class SimpleSecurityTransport(ISecureTransport):
 
 
 class SimpleSecureConn(ISecureConn):
-    def __init__(self, conn, key_phrase):
+    conn: "IRawConnection"
+    key_phrase: str
+    details: "TSecurityDetails"
+
+    def __init__(self, conn: "IRawConnection", key_phrase: str) -> None:
         self.conn = conn
-        self.details = {}
+        self.details = cast("TSecurityDetails", {})
         self.details["key_phrase"] = key_phrase
 
-    def get_conn(self):
+    def get_conn(self) -> "ISecureConn":
         """
         :return: connection object that has been made secure
         """
-        return self.conn
+        return cast("ISecureConn", self.conn)
 
-    def get_security_details(self):
+    def get_security_details(self) -> "TSecurityDetails":
         """
         :return: map containing details about the connections security
         """

libp2p/security/typing.py

@@ -0,0 +1,4 @@
+from typing import TypeVar, Dict, Any, NewType
+
+
+TSecurityDetails = NewType("TSecurityDetails", Dict[str, Any])