def test_rsa_signature_validator(): receiver_validator = RSASignatureValidator() sender_validator = RSASignatureValidator(RSAPrivateKey()) mallory_validator = RSASignatureValidator(RSAPrivateKey()) plain_record = DHTRecord(key=b'key', subkey=b'subkey', value=b'value', expiration_time=get_dht_time() + 10) protected_records = [ dataclasses.replace(plain_record, key=plain_record.key + sender_validator.local_public_key), dataclasses.replace(plain_record, subkey=plain_record.subkey + sender_validator.local_public_key), ] # test 1: Non-protected record (no signature added) assert sender_validator.sign_value(plain_record) == plain_record.value assert receiver_validator.validate(plain_record) # test 2: Correct signatures signed_records = [dataclasses.replace(record, value=sender_validator.sign_value(record)) for record in protected_records] for record in signed_records: assert receiver_validator.validate(record) assert receiver_validator.strip_value(record) == b'value' # test 3: Invalid signatures signed_records = protected_records # Without signature signed_records += [dataclasses.replace(record, value=record.value + b'[signature:INVALID_BYTES]') for record in protected_records] # With invalid signature signed_records += [dataclasses.replace(record, value=mallory_validator.sign_value(record)) for record in protected_records] # With someone else's signature for record in signed_records: assert not receiver_validator.validate(record)
async def test_auth_rpc_wrapper(): class Servicer: async def rpc_increment( self, request: dht_pb2.PingRequest) -> dht_pb2.PingResponse: assert request.peer.endpoint == '127.0.0.1:1111' assert request.auth.client_access_token.username == 'alice' response = dht_pb2.PingResponse() response.sender_endpoint = '127.0.0.1:2222' return response class Client: def __init__(self, servicer: Servicer): self._servicer = servicer async def rpc_increment( self, request: dht_pb2.PingRequest) -> dht_pb2.PingResponse: return await self._servicer.rpc_increment(request) servicer = AuthRPCWrapper(Servicer(), AuthRole.SERVICER, MockAuthorizer(RSAPrivateKey(), 'bob')) client = AuthRPCWrapper(Client(servicer), AuthRole.CLIENT, MockAuthorizer(RSAPrivateKey(), 'alice')) request = dht_pb2.PingRequest() request.peer.endpoint = '127.0.0.1:1111' response = await client.rpc_increment(request) assert response.sender_endpoint == '127.0.0.1:2222' assert response.auth.service_access_token.username == 'bob'
async def test_dhtnode_signatures(): alice = await hivemind.DHTNode.create(record_validator=RSASignatureValidator()) bob = await hivemind.DHTNode.create( record_validator=RSASignatureValidator(RSAPrivateKey()), initial_peers=[f"{LOCALHOST}:{alice.port}"]) mallory = await hivemind.DHTNode.create( record_validator=RSASignatureValidator(RSAPrivateKey()), initial_peers=[f"{LOCALHOST}:{alice.port}"]) key = b'key' subkey = b'protected_subkey' + bob.protocol.record_validator.local_public_key assert await bob.store(key, b'true_value', hivemind.get_dht_time() + 10, subkey=subkey) assert (await alice.get(key, latest=True)).value[subkey].value == b'true_value' store_ok = await mallory.store(key, b'fake_value', hivemind.get_dht_time() + 10, subkey=subkey) assert not store_ok assert (await alice.get(key, latest=True)).value[subkey].value == b'true_value' assert await bob.store(key, b'updated_true_value', hivemind.get_dht_time() + 10, subkey=subkey) assert (await alice.get(key, latest=True)).value[subkey].value == b'updated_true_value' await bob.shutdown() # Bob has shut down, now Mallory is the single peer of Alice store_ok = await mallory.store(key, b'updated_fake_value', hivemind.get_dht_time() + 10, subkey=subkey) assert not store_ok assert (await alice.get(key, latest=True)).value[subkey].value == b'updated_true_value'
async def test_valid_request_and_response(): client_authorizer = MockAuthorizer(RSAPrivateKey()) service_authorizer = MockAuthorizer(RSAPrivateKey()) request = dht_pb2.PingRequest() request.peer.endpoint = '127.0.0.1:7777' await client_authorizer.sign_request(request, service_authorizer.local_public_key) assert await service_authorizer.validate_request(request) response = dht_pb2.PingResponse() response.sender_endpoint = '127.0.0.1:31337' await service_authorizer.sign_response(response, request) assert await client_authorizer.validate_response(response, request)
def test_cached_key(): first_validator = RSASignatureValidator() second_validator = RSASignatureValidator() assert first_validator.local_public_key == second_validator.local_public_key third_validator = RSASignatureValidator(RSAPrivateKey()) assert first_validator.local_public_key != third_validator.local_public_key
def __init__(self, private_key: Optional[RSAPrivateKey] = None): if private_key is None: private_key = RSAPrivateKey.process_wide() self._private_key = private_key serialized_public_key = private_key.get_public_key().to_bytes() self._local_public_key = self.PUBLIC_KEY_FORMAT.replace( b'_key_', serialized_public_key)
def __init__(self, local_private_key: Optional[RSAPrivateKey]=None): if local_private_key is None: local_private_key = RSAPrivateKey.process_wide() self._local_private_key = local_private_key self._local_public_key = local_private_key.get_public_key() self._local_access_token = None self._refresh_lock = asyncio.Lock() self._recent_nonces = TimedStorage()
async def test_invalid_access_token(): client_authorizer = MockAuthorizer(RSAPrivateKey()) service_authorizer = MockAuthorizer(RSAPrivateKey()) request = dht_pb2.PingRequest() request.peer.endpoint = '127.0.0.1:7777' await client_authorizer.sign_request(request, service_authorizer.local_public_key) # Break the access token signature request.auth.client_access_token.signature = b'broken' assert not await service_authorizer.validate_request(request) response = dht_pb2.PingResponse() response.sender_endpoint = '127.0.0.1:31337' await service_authorizer.sign_response(response, request) # Break the access token signature response.auth.service_access_token.signature = b'broken' assert not await client_authorizer.validate_response(response, request)
async def test_invalid_signatures(): client_authorizer = MockAuthorizer(RSAPrivateKey()) service_authorizer = MockAuthorizer(RSAPrivateKey()) request = dht_pb2.PingRequest() request.peer.endpoint = '127.0.0.1:7777' await client_authorizer.sign_request(request, service_authorizer.local_public_key) # A man-in-the-middle attacker changes the request content request.peer.endpoint = '127.0.0.2:7777' assert not await service_authorizer.validate_request(request) response = dht_pb2.PingResponse() response.sender_endpoint = '127.0.0.1:31337' await service_authorizer.sign_response(response, request) # A man-in-the-middle attacker changes the response content response.sender_endpoint = '127.0.0.2:31337' assert not await client_authorizer.validate_response(response, request)
async def get_token(self) -> AccessToken: if MockAuthorizer._authority_private_key is None: MockAuthorizer._authority_private_key = RSAPrivateKey() self._authority_public_key = MockAuthorizer._authority_private_key.get_public_key( ) token = AccessToken(username=self._username, public_key=self.local_public_key.to_bytes(), expiration_time=str(datetime.utcnow() + timedelta(minutes=1))) token.signature = MockAuthorizer._authority_private_key.sign( self._token_to_bytes(token)) return token