Typing and linting fixes

This commit is contained in:
Alex Stokes 2019-08-23 23:43:36 +02:00
parent 0fa3331b8c
commit 1adef05e94
No known key found for this signature in database
GPG Key ID: 51CE1721B245C086
8 changed files with 53 additions and 33 deletions

View File

@ -98,7 +98,7 @@ def initialize_pair(
authenticator.update(tag) authenticator.update(tag)
tag = authenticator.digest() tag = authenticator.digest()
half = len(result) / 2 half = int(len(result) / 2)
first_half = result[:half] first_half = result[:half]
second_half = result[half:] second_half = result[half:]

View File

@ -1,3 +1,5 @@
from typing import cast
from Crypto.PublicKey import ECC from Crypto.PublicKey import ECC
from Crypto.PublicKey.ECC import EccKey from Crypto.PublicKey.ECC import EccKey
@ -9,7 +11,7 @@ class ECCPublicKey(PublicKey):
self.impl = impl self.impl = impl
def to_bytes(self) -> bytes: def to_bytes(self) -> bytes:
return self.impl.export_key("DER") return cast(bytes, self.impl.export_key(format="DER"))
@classmethod @classmethod
def from_bytes(cls, data: bytes) -> "ECCPublicKey": def from_bytes(cls, data: bytes) -> "ECCPublicKey":
@ -33,7 +35,7 @@ class ECCPrivateKey(PrivateKey):
return cls(private_key_impl) return cls(private_key_impl)
def to_bytes(self) -> bytes: def to_bytes(self) -> bytes:
return self.impl.export_key("DER") return cast(bytes, self.impl.export_key(format="DER"))
def get_type(self) -> KeyType: def get_type(self) -> KeyType:
return KeyType.ECC_P256 return KeyType.ECC_P256
@ -42,7 +44,7 @@ class ECCPrivateKey(PrivateKey):
raise NotImplementedError raise NotImplementedError
def get_public_key(self) -> PublicKey: def get_public_key(self) -> PublicKey:
return ECCPublicKey(self.impl.publickey()) return ECCPublicKey(self.impl.public_key())
def create_new_key_pair(curve: str) -> KeyPair: def create_new_key_pair(curve: str) -> KeyPair:

View File

@ -1,8 +1,8 @@
from typing import Callable, Tuple from typing import Callable, Tuple, cast
import Crypto.PublicKey.ECC as ECC import Crypto.PublicKey.ECC as ECC
from libp2p.crypto.ecc import create_new_key_pair from libp2p.crypto.ecc import ECCPrivateKey, create_new_key_pair
from libp2p.crypto.keys import PublicKey from libp2p.crypto.keys import PublicKey
SharedKeyGenerator = Callable[[bytes], bytes] SharedKeyGenerator = Callable[[bytes], bytes]
@ -20,7 +20,8 @@ def create_ephemeral_key_pair(curve_type: str) -> Tuple[PublicKey, SharedKeyGene
def _key_exchange(serialized_remote_public_key: bytes) -> bytes: def _key_exchange(serialized_remote_public_key: bytes) -> bytes:
remote_public_key = ECC.import_key(serialized_remote_public_key) remote_public_key = ECC.import_key(serialized_remote_public_key)
curve_point = remote_public_key.pointQ curve_point = remote_public_key.pointQ
secret_point = curve_point * key_pair.private_key.impl.d private_key = cast(ECCPrivateKey, key_pair.private_key)
secret_point = curve_point * private_key.impl.d
byte_size = secret_point.size_in_bytes() byte_size = secret_point.size_in_bytes()
return secret_point.x.to_bytes(byte_size, byteorder="big") return secret_point.x.to_bytes(byte_size, byteorder="big")

View File

@ -33,8 +33,10 @@ class Key(ABC):
""" """
... ...
def __eq__(self, other: "Key") -> bool: def __eq__(self, other: object) -> bool:
return self.impl == other.impl if not isinstance(other, Key):
return NotImplemented
return self.to_bytes() == other.to_bytes()
class PublicKey(Key): class PublicKey(Key):
@ -66,9 +68,7 @@ class PublicKey(Key):
@classmethod @classmethod
def deserialize_from_protobuf(cls, protobuf_data: bytes) -> protobuf.PublicKey: def deserialize_from_protobuf(cls, protobuf_data: bytes) -> protobuf.PublicKey:
protobuf_key = protobuf.PublicKey() return protobuf.PublicKey.FromString(protobuf_data)
protobuf_key.ParseFromString(protobuf_data)
return protobuf_key
class PrivateKey(Key): class PrivateKey(Key):
@ -110,9 +110,7 @@ class PrivateKey(Key):
@classmethod @classmethod
def deserialize_from_protobuf(cls, protobuf_data: bytes) -> protobuf.PrivateKey: def deserialize_from_protobuf(cls, protobuf_data: bytes) -> protobuf.PrivateKey:
protobuf_key = protobuf.PrivateKey() return protobuf.PrivateKey.FromString(protobuf_data)
protobuf_key.ParseFromString(protobuf_data)
return protobuf_key
@dataclass(frozen=True) @dataclass(frozen=True)

View File

@ -2,7 +2,7 @@ import asyncio
SIZE_LEN_BYTES = 4 SIZE_LEN_BYTES = 4
# TODO unify w/ https://github.com/libp2p/py-libp2p/blob/1aed52856f56a4b791696bbcbac31b5f9c2e88c9/libp2p/utils.py#L85-L99 # TODO unify w/ https://github.com/libp2p/py-libp2p/blob/1aed52856f56a4b791696bbcbac31b5f9c2e88c9/libp2p/utils.py#L85-L99 # noqa: E501
def encode(msg_bytes: bytes) -> bytes: def encode(msg_bytes: bytes) -> bytes:

View File

@ -3,7 +3,6 @@ from typing import Optional
from libp2p.crypto.keys import PrivateKey, PublicKey from libp2p.crypto.keys import PrivateKey, PublicKey
from libp2p.network.connection.raw_connection_interface import IRawConnection from libp2p.network.connection.raw_connection_interface import IRawConnection
from libp2p.peer.id import ID from libp2p.peer.id import ID
from libp2p.security.base_transport import BaseSecureTransport
from libp2p.security.secure_conn_interface import ISecureConn from libp2p.security.secure_conn_interface import ISecureConn
@ -21,12 +20,13 @@ class BaseSession(ISecureConn):
def __init__( def __init__(
self, self,
transport: BaseSecureTransport, local_peer: ID,
local_private_key: PrivateKey,
conn: IRawConnection, conn: IRawConnection,
peer_id: Optional[ID] = None, peer_id: Optional[ID] = None,
) -> None: ) -> None:
self.local_peer = transport.local_peer self.local_peer = local_peer
self.local_private_key = transport.local_private_key self.local_private_key = local_private_key
self.remote_peer_id = peer_id self.remote_peer_id = peer_id
self.remote_permanent_pubkey = None self.remote_permanent_pubkey = None

View File

@ -21,3 +21,7 @@ class InvalidSignatureOnExchange(SecioException):
class HandshakeFailed(SecioException): class HandshakeFailed(SecioException):
pass pass
class IncompatibleChoices(SecioException):
pass

View File

@ -25,6 +25,7 @@ from libp2p.security.secure_conn_interface import ISecureConn
from .exceptions import ( from .exceptions import (
HandshakeFailed, HandshakeFailed,
IncompatibleChoices,
InvalidSignatureOnExchange, InvalidSignatureOnExchange,
PeerMismatchException, PeerMismatchException,
SecioException, SecioException,
@ -43,17 +44,20 @@ DEFAULT_SUPPORTED_CIPHERS = "AES-128"
DEFAULT_SUPPORTED_HASHES = "SHA256" DEFAULT_SUPPORTED_HASHES = "SHA256"
@dataclass
class SecureSession(BaseSession): class SecureSession(BaseSession):
local_peer: PeerID def __init__(
local_encryption_parameters: AuthenticatedEncryptionParameters self,
local_peer: PeerID,
local_private_key: PrivateKey,
local_encryption_parameters: AuthenticatedEncryptionParameters,
remote_peer: PeerID,
remote_encryption_parameters: AuthenticatedEncryptionParameters,
conn: IRawConnection,
) -> None:
super().__init__(local_peer, local_private_key, conn, remote_peer)
remote_peer: PeerID self.local_encryption_parameters = local_encryption_parameters
remote_encryption_parameters: AuthenticatedEncryptionParameters self.remote_encryption_parameters = remote_encryption_parameters
conn: IRawConnection
def __post_init__(self):
self._initialize_authenticated_encryption_for_local_peer() self._initialize_authenticated_encryption_for_local_peer()
self._initialize_authenticated_encryption_for_remote_peer() self._initialize_authenticated_encryption_for_remote_peer()
@ -68,7 +72,8 @@ class SecureSession(BaseSession):
async def _read_msg(self) -> bytes: async def _read_msg(self) -> bytes:
# TODO do we need to serialize reads? # TODO do we need to serialize reads?
msg = await read_next_message(self.conn) # TODO do not expose reader
msg = await read_next_message(self.conn.reader)
return self.remote_encrypter.decrypt_if_valid(msg) return self.remote_encrypter.decrypt_if_valid(msg)
async def write(self, data: bytes) -> None: async def write(self, data: bytes) -> None:
@ -135,6 +140,9 @@ class EncryptionParameters:
ephemeral_public_key: PublicKey ephemeral_public_key: PublicKey
def __init__(self) -> None:
pass
@dataclass @dataclass
class SessionParameters: class SessionParameters:
@ -148,6 +156,9 @@ class SessionParameters:
order: int order: int
shared_key: bytes shared_key: bytes
def __init__(self) -> None:
pass
async def _response_to_msg(conn: IRawConnection, msg: bytes) -> bytes: async def _response_to_msg(conn: IRawConnection, msg: bytes) -> bytes:
# TODO clean up ``IRawConnection`` so that we don't have to break # TODO clean up ``IRawConnection`` so that we don't have to break
@ -182,6 +193,7 @@ def _select_parameter_from_order(
for first, second in zip(first_choices, second_choices): for first, second in zip(first_choices, second_choices):
if first == second: if first == second:
return first return first
raise IncompatibleChoices()
def _select_encryption_parameters( def _select_encryption_parameters(
@ -302,7 +314,9 @@ async def _establish_session_parameters(
def _mk_session_from( def _mk_session_from(
session_parameters: SessionParameters, conn: IRawConnection local_private_key: PrivateKey,
session_parameters: SessionParameters,
conn: IRawConnection,
) -> SecureSession: ) -> SecureSession:
key_set1, key_set2 = initialize_pair_for_encryption( key_set1, key_set2 = initialize_pair_for_encryption(
session_parameters.local_encryption_parameters.cipher_type, session_parameters.local_encryption_parameters.cipher_type,
@ -315,6 +329,7 @@ def _mk_session_from(
session = SecureSession( session = SecureSession(
session_parameters.local_peer, session_parameters.local_peer,
local_private_key,
key_set1, key_set1,
session_parameters.remote_peer, session_parameters.remote_peer,
key_set2, key_set2,
@ -329,7 +344,7 @@ async def _finish_handshake(session: ISecureConn, remote_nonce: bytes) -> bytes:
async def create_secure_session( async def create_secure_session(
transport: BaseSecureTransport, conn: IRawConnection, remote_peer: PeerID = None transport: "SecIOTransport", conn: IRawConnection, remote_peer: PeerID = None
) -> ISecureConn: ) -> ISecureConn:
""" """
Attempt the initial `secio` handshake with the remote peer. Attempt the initial `secio` handshake with the remote peer.
@ -348,7 +363,7 @@ async def create_secure_session(
conn.close() conn.close()
raise e raise e
session = _mk_session_from(session_parameters, conn) session = _mk_session_from(local_private_key, session_parameters, conn)
received_nonce = await _finish_handshake(session, remote_nonce) received_nonce = await _finish_handshake(session, remote_nonce)
if received_nonce != local_nonce: if received_nonce != local_nonce: