Allow registration of alias

snitun/server/peer.py

@@ -4,7 +4,7 @@
 import hashlib
 import logging
 import os
-from typing import Optional, Coroutine
+from typing import Optional, Coroutine, List
 
 import async_timeout
 
@@ -25,11 +25,13 @@ def __init__(
         aes_key: bytes,
         aes_iv: bytes,
         throttling: Optional[int] = None,
+        alias: Optional[List[str]] = None,
     ):
         """Initialize a Peer."""
         self._hostname = hostname
         self._valid = valid
         self._throttling = throttling
+        self._alias = alias or []
         self._multiplexer = None
         self._crypto = CryptoTransport(aes_key, aes_iv)
 
@@ -38,6 +40,11 @@ def hostname(self) -> str:
         """Return his hostname."""
         return self._hostname
 
+    @property
+    def alias(self) -> List[str]:
+        """Return the alias."""
+        return self._alias
+
     @property
     def is_connected(self) -> bool:
         """Return True if we are connected to peer."""

snitun/server/peer_manager.py

@@ -60,7 +60,14 @@ def create_peer(self, fernet_data: bytes) -> Peer:
         aes_key = bytes.fromhex(config["aes_key"])
         aes_iv = bytes.fromhex(config["aes_iv"])
 
-        return Peer(hostname, valid, aes_key, aes_iv, throttling=self._throttling)
+        return Peer(
+            hostname,
+            valid,
+            aes_key,
+            aes_iv,
+            throttling=self._throttling,
+            alias=config.get("alias", []),
+        )
 
     def add_peer(self, peer: Peer) -> None:
         """Register peer to internal hostname list."""
@@ -70,6 +77,9 @@ def add_peer(self, peer: Peer) -> None:
 
         _LOGGER.debug("New peer connection: %s", peer.hostname)
         self._peers[peer.hostname] = peer
+        for alias in peer.alias:
+            _LOGGER.debug("New peer connection alias: %s for %s", alias, peer.hostname)
+            self._peers[alias] = peer
 
         if self._event_callback:
             self._loop.call_soon(self._event_callback, peer, PeerManagerEvent.CONNECTED)
@@ -80,6 +90,8 @@ def remove_peer(self, peer: Peer) -> None:
             return
         _LOGGER.debug("Close peer connection: %s", peer.hostname)
         self._peers.pop(peer.hostname)
+        for alias in peer.alias:
+            self._peers.pop(alias, None)
 
         if self._event_callback:
             self._loop.call_soon(

tests/client/test_client_peer.py

@@ -41,6 +41,37 @@ async def test_init_client_peer(peer_listener, peer_manager, test_endpoint):
     assert not peer_manager.peer_available("localhost")
 
 
+async def test_init_client_peer_with_alias(peer_listener, peer_manager, test_endpoint):
+    """Test setup of ClientPeer with custom tomain."""
+    client = ClientPeer("127.0.0.1", "8893")
+    connector = Connector("127.0.0.1", "8822")
+
+    assert not client.is_connected
+    assert not peer_manager.peer_available("localhost")
+    assert not peer_manager.peer_available("localhost.custom")
+
+    valid = datetime.utcnow() + timedelta(days=1)
+    aes_key = os.urandom(32)
+    aes_iv = os.urandom(16)
+    hostname = "localhost"
+    fernet_token = create_peer_config(
+        valid.timestamp(), hostname, aes_key, aes_iv, alias=["localhost.custom"]
+    )
+
+    await client.start(connector, fernet_token, aes_key, aes_iv)
+    await asyncio.sleep(0.1)
+    assert peer_manager.peer_available("localhost")
+    assert peer_manager.peer_available("localhost.custom")
+    assert client.is_connected
+    assert client._multiplexer._throttling is None
+
+    await client.stop()
+    await asyncio.sleep(0.1)
+    assert not client.is_connected
+    assert not peer_manager.peer_available("localhost")
+    assert not peer_manager.peer_available("localhost.custom")
+
+
 async def test_init_client_peer_invalid_token(
     peer_listener, peer_manager, test_endpoint
 ):

tests/server/const_fernet.py

@@ -1,5 +1,6 @@
 """Const value for Fernet tests."""
 import json
+from typing import List, Optional
 
 from cryptography.fernet import Fernet, MultiFernet
 
@@ -10,7 +11,11 @@
 
 
 def create_peer_config(
-    valid: int, hostname: str, aes_key: bytes, aes_iv: bytes
+    valid: int,
+    hostname: str,
+    aes_key: bytes,
+    aes_iv: bytes,
+    alias: Optional[List[str]] = None,
 ) -> bytes:
     """Create a fernet token."""
     fernet = MultiFernet([Fernet(key) for key in FERNET_TOKENS])
@@ -20,6 +25,7 @@ def create_peer_config(
             {
                 "valid": valid,
                 "hostname": hostname,
+                "alias": alias or [],
                 "aes_key": aes_key.hex(),
                 "aes_iv": aes_iv.hex(),
             }

tests/server/test_peer.py

@@ -15,11 +15,18 @@
 def test_init_peer():
     """Test simple init of peer."""
     valid = datetime.utcnow() + timedelta(days=1)
-    peer = Peer("localhost", valid, os.urandom(32), os.urandom(16))
+    peer = Peer(
+        "localhost",
+        valid,
+        os.urandom(32),
+        os.urandom(16),
+        alias="localhost.custom",
+    )
 
     assert peer.is_valid
     assert peer.hostname == "localhost"
     assert peer.multiplexer is None
+    assert peer.alias == "localhost.custom"
 
 
 async def test_init_peer_multiplexer(event_loop, test_client, test_server):

tests/server/test_peer_manager.py

@@ -47,6 +47,41 @@ async def test_init_new_peer():
     assert manager.connections == 1
 
 
+async def test_init_new_peer_with_alias():
+    """Init a new peer with custom domain."""
+    manager = PeerManager(FERNET_TOKENS)
+
+    valid = datetime.utcnow() + timedelta(days=1)
+    aes_key = os.urandom(32)
+    aes_iv = os.urandom(16)
+    hostname = "localhost"
+    alias = "localhost.custom"
+    fernet_token = create_peer_config(
+        valid.timestamp(), hostname, aes_key, aes_iv, alias=[alias]
+    )
+
+    peer = manager.create_peer(fernet_token)
+    assert peer.hostname == hostname
+    assert peer.alias == [alias]
+    assert not peer.is_ready
+    assert not manager.get_peer(hostname)
+    assert not manager.get_peer(alias)
+    assert not manager.peer_available(hostname)
+    assert not manager.peer_available(alias)
+    assert hostname not in manager._peers
+    assert alias not in manager._peers
+    assert manager.connections == 0
+
+    manager.add_peer(peer)
+    assert manager.get_peer(hostname)
+    assert manager.get_peer(alias)
+    assert not manager.peer_available(hostname)
+    assert not manager.peer_available(alias)
+    assert hostname in manager._peers
+    assert alias in manager._peers
+    assert manager.connections == 2
+
+
 async def test_init_new_peer_not_valid_time():
     """Init a new peer."""
     manager = PeerManager(FERNET_TOKENS)