Add NIP-42 authentication to relay discovery (#10)

* Add NIP-42 authentication to relay discovery

* Move discovery tests into tests directory

README.md

@@ -14,6 +14,9 @@ The repository contains three scripts that form a simple pipeline:
   - Starts from a seed relay URL and connects to it via WebSocket (ws/wss).
   - Breadth‑first searches follow lists (kind 3) and relay lists (kind 10002), extracting relay URLs from event tags (r/p tags).
   - Tests each candidate relay for basic Nostr protocol responsiveness by issuing a REQ and evaluating responses (EVENT/EOSE/NOTICE).
+  - Responds to NIP-42 authentication challenges. An ephemeral Nostr key is
+    generated by default; set `NOSTR_PRIVATE_KEY` or pass `--private-key` with a
+    64-character hex secret when testing membership-restricted relays.
   - Processes relays in concurrent batches and periodically saves progress to `relay_discovery_results.json`.
 - Output: `relay_discovery_results.json` with:
   - `functioning_relays` (array of ws/wss relay URLs)
@@ -42,6 +45,8 @@ Note: The geolocation step is IPv4‑only (the DB‑IP file used here is IPv4).
 
 ## Tips and caveats
 - The discovery step uses timeouts and concurrent batches; tune `--batch-size`, `--timeout`, and `--max-depth` for your environment.
+- NIP-42 authentication proves control of a key but does not bypass relay
+  membership or payment policies. Use a relay-authorized key when required.
 - The BitChat filter posts a test kind 20000 event. Ensure `nak` is configured to publish (e.g., via its config or environment variables as per `nak` docs).
 - Geolocation is an estimate based on DB‑IP and only for IPv4. Accuracy varies and may reflect the ISP/hosting POP rather than precise server location.
 - `relays_geo_lookup.py` stops at the first IPv4 for which DB‑IP provides coordinates. Some hostnames resolve to multiple IPs.

nostr_relay_discovery.py

@@ -7,8 +7,10 @@ follow lists (kind 3 events) and analyzing relay tags to build a network map.
 """
 
 import asyncio
+import hashlib
 import json
 import logging
+import os
 import re
 import time
 import argparse
@@ -20,6 +22,7 @@ from typing import Set, List, Dict, Optional, Tuple
 from urllib.parse import urlparse
 import websockets
 import websockets.exceptions
+from coincurve import PrivateKey
 
 
 # Configure logging
@@ -60,13 +63,14 @@ class RelayDiscoveryStats:
 class NostrRelayDiscovery:
     """Nostr relay discovery tool using breadth-first search through follow lists"""
     
-    def __init__(self, initial_relay: str, max_depth: int = 3, connection_timeout: int = 5, output_file: str = "relay_discovery_results.json", save_point: int = SAVE_POINT, batch_size: int = 10):
+    def __init__(self, initial_relay: str, max_depth: int = 3, connection_timeout: int = 5, output_file: str = "relay_discovery_results.json", save_point: int = SAVE_POINT, batch_size: int = 10, private_key: Optional[str] = None):
         self.initial_relay = initial_relay
         self.max_depth = max_depth
         self.connection_timeout = connection_timeout
         self.output_file = output_file
         self.save_point = save_point
         self.batch_size = batch_size
+        self.private_key = self._load_private_key(private_key)
         
         # Discovery state
         self.to_visit: deque = deque()  # (relay_url, depth) - will be populated by load_existing_results
@@ -122,6 +126,119 @@ class NostrRelayDiscovery:
         random_bytes = secrets.token_bytes(16)
         return base64.b64encode(random_bytes).decode()
 
+    @staticmethod
+    def _load_private_key(private_key_hex: Optional[str]) -> PrivateKey:
+        """Load a configured Nostr secret or generate an ephemeral discovery key."""
+        if private_key_hex is None:
+            return PrivateKey()
+
+        try:
+            secret = bytes.fromhex(private_key_hex)
+        except ValueError as exc:
+            raise ValueError("--private-key must be 64 hexadecimal characters") from exc
+
+        if len(secret) != 32:
+            raise ValueError("--private-key must be 64 hexadecimal characters")
+
+        try:
+            return PrivateKey(secret)
+        except ValueError as exc:
+            raise ValueError("--private-key is not a valid secp256k1 secret") from exc
+
+    def build_auth_event(self, relay_url: str, challenge: str) -> Dict:
+        """Build and sign a NIP-42 kind 22242 authentication event."""
+        pubkey = self.private_key.public_key.format(compressed=True)[1:].hex()
+        event = {
+            "pubkey": pubkey,
+            "created_at": int(time.time()),
+            "kind": 22242,
+            "tags": [
+                ["relay", relay_url],
+                ["challenge", challenge],
+            ],
+            "content": "",
+        }
+        serialized = json.dumps(
+            [
+                0,
+                event["pubkey"],
+                event["created_at"],
+                event["kind"],
+                event["tags"],
+                event["content"],
+            ],
+            separators=(",", ":"),
+            ensure_ascii=False,
+        ).encode()
+        event_id = hashlib.sha256(serialized).digest()
+        event["id"] = event_id.hex()
+        event["sig"] = self.private_key.sign_schnorr(event_id).hex()
+        return event
+
+    async def receive_with_auth(
+        self, websocket, relay_url: str, timeout: float, request_message=None
+    ):
+        """Receive one application message, completing NIP-42 if challenged.
+
+        Some relays reject the request that triggered authentication before
+        processing the AUTH response. When ``request_message`` is supplied, it
+        is sent again after authentication succeeds.
+        """
+        deadline = asyncio.get_running_loop().time() + timeout
+
+        while True:
+            remaining = deadline - asyncio.get_running_loop().time()
+            if remaining <= 0:
+                raise asyncio.TimeoutError
+
+            response = await asyncio.wait_for(websocket.recv(), timeout=remaining)
+            data = json.loads(response)
+
+            if not isinstance(data, list) or not data:
+                return data
+
+            if data[0] != "AUTH":
+                return data
+
+            if len(data) < 2 or not isinstance(data[1], str):
+                raise ValueError("relay sent a malformed NIP-42 AUTH challenge")
+
+            auth_event = self.build_auth_event(relay_url, data[1])
+            await websocket.send(json.dumps(["AUTH", auth_event]))
+            logger.debug(f"Sent NIP-42 AUTH response to {relay_url}")
+
+            remaining = deadline - asyncio.get_running_loop().time()
+            if remaining <= 0:
+                raise asyncio.TimeoutError
+            while True:
+                remaining = deadline - asyncio.get_running_loop().time()
+                if remaining <= 0:
+                    raise asyncio.TimeoutError
+                auth_response = json.loads(
+                    await asyncio.wait_for(websocket.recv(), timeout=remaining)
+                )
+                if (
+                    isinstance(auth_response, list)
+                    and len(auth_response) >= 3
+                    and auth_response[0] == "OK"
+                    and auth_response[1] == auth_event["id"]
+                ):
+                    if auth_response[2] is not True:
+                        reason = (
+                            auth_response[3]
+                            if len(auth_response) > 3
+                            else "authentication rejected"
+                        )
+                        raise PermissionError(
+                            f"NIP-42 authentication failed: {reason}"
+                        )
+                    break
+
+            logger.debug(f"NIP-42 authentication succeeded with {relay_url}")
+            if request_message is not None:
+                await websocket.send(json.dumps(request_message))
+                logger.debug(f"Resent request after NIP-42 authentication: {relay_url}")
+    
     async def test_relay_connection(self, relay_url: str) -> bool:
         """Test if a relay is functioning by attempting to connect and validate Nostr protocol responses"""
         try:
@@ -146,12 +263,12 @@ class NostrRelayDiscovery:
                 
                 # Wait for a response and validate it's a proper Nostr protocol message
                 try:
-                    response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
-                    logger.debug(f"Received response: {response[:200]}...")  # Log first 200 chars
+                    data = await self.receive_with_auth(
+                        websocket, relay_url, 5.0, req_msg
+                    )
+                    logger.debug(f"Received response: {str(data)[:200]}...")
                     
                     try:
-                        data = json.loads(response)
-                        
                         # Check if it's a valid Nostr protocol message
                         if not isinstance(data, list) or len(data) < 2:
                             logger.debug(f"Invalid Nostr message format from {relay_url}: not a list or too short")
@@ -186,7 +303,7 @@ class NostrRelayDiscovery:
                             logger.debug(f"Unexpected message type from {relay_url}: {message_type}")
                             return False
                             
-                    except json.JSONDecodeError as e:
+                    except (json.JSONDecodeError, ValueError) as e:
                         logger.debug(f"Invalid JSON response from {relay_url}: {e}")
                         return False
                         
@@ -261,8 +378,9 @@ class NostrRelayDiscovery:
                 
                 while time.time() - start_time < timeout_duration:
                     try:
-                        message = await asyncio.wait_for(websocket.recv(), timeout=5.0)
-                        data = json.loads(message)
+                        data = await self.receive_with_auth(
+                            websocket, relay_url, 5.0, req_msg
+                        )
                         
                         if data[0] == "EVENT" and data[1] == subscription_id:
                             event = data[2]
@@ -543,6 +661,14 @@ async def main():
         default=10,
         help="Number of relays to process concurrently (default: 10)"
     )
+    parser.add_argument(
+        "--private-key",
+        default=os.environ.get("NOSTR_PRIVATE_KEY"),
+        help=(
+            "64-character hex Nostr private key for NIP-42 authentication "
+            "(default: NOSTR_PRIVATE_KEY or an ephemeral key)"
+        ),
+    )
     parser.add_argument(
         "--verbose",
         action="store_true",
@@ -561,7 +687,8 @@ async def main():
         args.timeout, 
         args.output,
         args.save_point,
-        args.batch_size
+        args.batch_size,
+        args.private_key,
     )
     if not discovery.is_valid_relay_url(args.initial_relay):
         print(f"Error: Invalid relay URL: {args.initial_relay}")

requirements.txt

@@ -1,4 +1,5 @@
 websockets>=12.0
+coincurve>=20.0.0
 matplotlib>=3.5.0
 pandas>=1.4.0
 numpy>=1.20.0

tests/test_nostr_relay_discovery.py

@@ -0,0 +1,121 @@
+import asyncio
+import hashlib
+import json
+import unittest
+from unittest.mock import patch
+
+from coincurve import PublicKeyXOnly
+
+from nostr_relay_discovery import NostrRelayDiscovery
+
+
+class FakeWebSocket:
+    def __init__(self, messages):
+        self.messages = iter(messages)
+        self.sent = []
+
+    async def recv(self):
+        return next(self.messages)
+
+    async def send(self, message):
+        self.sent.append(json.loads(message))
+
+
+class Nip42Tests(unittest.TestCase):
+    def setUp(self):
+        self.discovery = NostrRelayDiscovery(
+            "wss://relay.example.com",
+            private_key="01".zfill(64),
+        )
+
+    def test_build_auth_event_is_valid_schnorr_event(self):
+        with patch("nostr_relay_discovery.time.time", return_value=1_700_000_000):
+            event = self.discovery.build_auth_event(
+                "wss://relay.example.com", "challenge"
+            )
+
+        serialized = json.dumps(
+            [
+                0,
+                event["pubkey"],
+                event["created_at"],
+                event["kind"],
+                event["tags"],
+                event["content"],
+            ],
+            separators=(",", ":"),
+            ensure_ascii=False,
+        ).encode()
+        event_id = hashlib.sha256(serialized).digest()
+
+        self.assertEqual(event["id"], event_id.hex())
+        self.assertTrue(
+            PublicKeyXOnly(bytes.fromhex(event["pubkey"])).verify(
+                bytes.fromhex(event["sig"]), event_id
+            )
+        )
+
+    def test_receive_with_auth_returns_subscription_response(self):
+        async def run_test():
+            auth_event = self.discovery.build_auth_event(
+                "wss://relay.example.com", "challenge"
+            )
+            websocket = FakeWebSocket(
+                [
+                    json.dumps(["AUTH", "challenge"]),
+                    json.dumps(
+                        ["CLOSED", "subscription", "auth-required: authenticate"]
+                    ),
+                    json.dumps(["OK", auth_event["id"], True, ""]),
+                    json.dumps(["EOSE", "subscription"]),
+                ]
+            )
+
+            with patch.object(
+                self.discovery, "build_auth_event", return_value=auth_event
+            ):
+                response = await self.discovery.receive_with_auth(
+                    websocket,
+                    "wss://relay.example.com",
+                    5,
+                    ["REQ", "subscription", {"kinds": [1]}],
+                )
+
+            self.assertEqual(response, ["EOSE", "subscription"])
+            self.assertEqual(
+                websocket.sent,
+                [
+                    ["AUTH", auth_event],
+                    ["REQ", "subscription", {"kinds": [1]}],
+                ],
+            )
+
+        asyncio.run(run_test())
+
+    def test_receive_with_auth_reports_rejection(self):
+        async def run_test():
+            auth_event = self.discovery.build_auth_event(
+                "wss://relay.example.com", "challenge"
+            )
+            websocket = FakeWebSocket(
+                [
+                    json.dumps(["AUTH", "challenge"]),
+                    json.dumps(
+                        ["OK", auth_event["id"], False, "restricted: not a member"]
+                    ),
+                ]
+            )
+
+            with patch.object(
+                self.discovery, "build_auth_event", return_value=auth_event
+            ):
+                with self.assertRaisesRegex(PermissionError, "not a member"):
+                    await self.discovery.receive_with_auth(
+                        websocket, "wss://relay.example.com", 5
+                    )
+
+        asyncio.run(run_test())
+
+
+if __name__ == "__main__":
+    unittest.main()