Merge pull request #7 from kiraum/kiraum/issues_2

feat: add option to specify the base url other than irrexplorer.net

irrexplorer_cli/helpers.py

@@ -18,12 +18,14 @@ def validate_prefix_format(prefix_input: str) -> bool:
 
 def validate_asn_format(asn_input: str) -> bool:
     """Validate ASN format."""
+    if not isinstance(asn_input, str):
+        return False
     asn_pattern = r"^(?:AS|as)?(\d+)$"
     match = re.match(asn_pattern, asn_input)
-    if match:
-        asn_number = int(match.group(1))
-        return 0 <= asn_number <= 4294967295
-    return False
+    if not match:
+        return False
+    asn_number = int(match.group(1))
+    return 0 <= asn_number <= 4294967295
 
 
 def format_prefix_result(result: PrefixInfo, prefix_type: str) -> str:
@@ -121,3 +123,9 @@ async def find_least_specific_prefix(direct_overlaps: List[PrefixInfo]) -> str |
             if least_specific is None or int(mask) < int(least_specific.split("/")[1]):
                 least_specific = info.prefix
     return least_specific
+
+
+def validate_url_format(url: str) -> bool:
+    """Validate URL format."""
+    url_pattern = r"^https?://[a-zA-Z0-9.-]+(?:\.[a-zA-Z]{2,})+(?:/[^\s]*)?$"
+    return bool(re.match(url_pattern, url))

irrexplorer_cli/main.py

@@ -2,12 +2,12 @@
 
 import asyncio
 from importlib.metadata import version
+from typing import Optional
 
 import typer
-from click import Context
 from rich.console import Console
 
-from irrexplorer_cli.helpers import validate_asn_format, validate_prefix_format
+from irrexplorer_cli.helpers import validate_asn_format, validate_prefix_format, validate_url_format
 from irrexplorer_cli.queries import async_asn_query, async_prefix_query
 
 __version__ = version("irrexplorer-cli")
@@ -30,18 +30,23 @@ def version_display(display_version: bool) -> None:
 
 @app.callback()
 def callback(
+    ctx: typer.Context,
     _: bool = typer.Option(None, "--version", "-v", callback=version_display, is_eager=True),
+    base_url: Optional[str] = typer.Option(None, "--url", "-u", help="Base URL for IRR Explorer API"),
 ) -> None:
     """Query IRR Explorer for prefix information."""
+    ctx.ensure_object(dict)
+    ctx.obj["base_url"] = base_url
 
 
 @app.command(no_args_is_help=True)
 def prefix(
+    ctx: typer.Context,
     prefix_query: str = typer.Argument(None, help="Prefix to query (e.g., 193.0.0.0/21)"),
-    output_format: str = typer.Option(None, "--format", "-f", help="Output format (json or csv)"),
-    ctx: Context = CTX_OPTION,
+    output_format: Optional[str] = typer.Option(None, "--format", "-f", help="Output format (json or csv)"),
 ) -> None:
     """Query IRR Explorer for prefix information."""
+    base_url: Optional[str] = ctx.obj.get("base_url")
     if not prefix_query:
         if ctx:
             typer.echo(ctx.get_help())
@@ -51,28 +56,38 @@ def prefix(
         typer.echo(f"Error: Invalid prefix format: {prefix_query}")
         raise typer.Exit(1)
 
-    asyncio.run(async_prefix_query(prefix_query, output_format))
+    if base_url and not validate_url_format(base_url):
+        typer.echo(f"Error: Invalid URL format: {base_url}")
+        raise typer.Exit(1)
+
+    asyncio.run(async_prefix_query(prefix_query, output_format, base_url))
 
 
-@app.command()
+@app.command(no_args_is_help=True)
 def asn(
+    ctx: typer.Context,
     asn_query: str = typer.Argument(None, help="AS number to query (e.g., AS2111, as2111, or 2111)"),
-    output_format: str = typer.Option(None, "--format", "-f", help="Output format (json or csv)"),
-    ctx: Context = CTX_OPTION,
+    output_format: Optional[str] = typer.Option(None, "--format", "-f", help="Output format (json or csv)"),
 ) -> None:
     """Query IRR Explorer for AS number information."""
+    base_url: Optional[str] = ctx.obj.get("base_url")
     if not asn_query:
         if ctx:
             typer.echo(ctx.get_help())
         raise typer.Exit()
 
+    if isinstance(asn_query, str):
+        if not asn_query.upper().startswith("AS"):
+            asn_query = f"AS{asn_query}"
+        else:
+            asn_query = f"AS{asn_query[2:]}"
+
     if not validate_asn_format(asn_query):
         typer.echo(f"Error: Invalid ASN format: {asn_query}")
         raise typer.Exit(1)
 
-    if not asn_query.upper().startswith("AS"):
-        asn_query = f"AS{asn_query}"
-    else:
-        asn_query = f"AS{asn_query[2:]}"
+    if base_url and not validate_url_format(base_url):
+        typer.echo(f"Error: Invalid URL format: {base_url}")
+        raise typer.Exit(1)
 
-    asyncio.run(async_asn_query(asn_query, output_format))
+    asyncio.run(async_asn_query(asn_query, output_format, base_url))

irrexplorer_cli/queries.py

@@ -4,6 +4,7 @@
 from typing import Optional
 
 import httpx
+import typer
 
 from irrexplorer_cli.helpers import (
     find_least_specific_prefix,
@@ -25,35 +26,40 @@ async def process_overlaps(explorer: IrrExplorer, least_specific: str) -> None:
         pass
 
 
-async def async_prefix_query(pfx: str, output_format: Optional[str] = None) -> None:
+async def async_prefix_query(pfx: str, output_format: Optional[str] = None, base_url: Optional[str] = None) -> None:
     """Execute asynchronous prefix query and display results."""
-    explorer = IrrExplorer()
+    explorer = IrrExplorer(base_url=base_url) if base_url else IrrExplorer()
     display = IrrDisplay()
+
     try:
         direct_overlaps = await explorer.fetch_prefix_info(pfx)
-
         if output_format == "json":
             json_data = [result.model_dump() for result in direct_overlaps]
             print(json.dumps(json_data, indent=2))
         elif output_format == "csv":
             print("Type,Prefix,Category,RIR,RPKI_Status,BGP_Origins,IRR_Routes,Messages")
-
             for result in direct_overlaps:
                 print(format_prefix_result(result, "DIRECT"))
-
             least_specific = await find_least_specific_prefix(direct_overlaps)
             if least_specific:
                 await process_overlaps(explorer, least_specific)
         else:
             await display.display_prefix_info(direct_overlaps)
+    except httpx.ConnectError as exc:
+        print(
+            f"Error: Unable to connect to {base_url or 'default IRR Explorer instance'}. "
+            "Please verify the URL is correct and the service is available."
+        )
+        raise typer.Exit(1) from exc
     finally:
         await explorer.close()
 
 
-async def async_asn_query(as_number: str, output_format: Optional[str] = None) -> None:
+async def async_asn_query(as_number: str, output_format: Optional[str] = None, base_url: Optional[str] = None) -> None:
     """Execute asynchronous ASN query and display results."""
-    explorer = IrrExplorer()
+    explorer = IrrExplorer(base_url=base_url) if base_url else IrrExplorer()
     display = IrrDisplay()
+
     try:
         results = await explorer.fetch_asn_info(as_number)
         sets_data = await explorer.fetch_asn_sets(as_number)
@@ -69,5 +75,11 @@ async def async_asn_query(as_number: str, output_format: Optional[str] = None) -
             print()
         else:
             await display.display_asn_info(results, as_number, sets_data)
+    except httpx.ConnectError as exc:
+        print(
+            f"Error: Unable to connect to {base_url or 'default IRR Explorer instance'}. "
+            "Please verify the URL is correct and the service is available."
+        )
+        raise typer.Exit(1) from exc
     finally:
         await explorer.close()

tests/test_cli.py

@@ -108,4 +108,4 @@ def test_asn_query_command() -> None:
     with patch("irrexplorer_cli.main.async_asn_query", return_value=None) as mock_query:
         result = runner.invoke(app, ["asn", "AS202196"])
         assert not result.exit_code
-        mock_query.assert_called_once_with("AS202196", None)
+        mock_query.assert_called_once_with("AS202196", None, None)

tests/test_helpers.py

@@ -12,6 +12,7 @@
     format_prefix_result,
     validate_asn_format,
     validate_prefix_format,
+    validate_url_format,
 )
 from tests.fixtures import (
     COMMON_ASN_DATA,
@@ -142,3 +143,18 @@ def test_format_as_sets_empty() -> None:
 def test_format_direct_origins_with_rpki_routes() -> None:
     """Test direct origins formatting with RPKI routes."""
     format_direct_origins("AS12345", COMMON_ASN_DATA)
+
+
+def test_validate_url_format() -> None:
+    """Test URL format validation."""
+    # Valid URLs
+    assert validate_url_format("https://example.com") is True
+    assert validate_url_format("http://sub.example.com/path") is True
+    assert validate_url_format("https://example.com/path?query=value") is True
+    assert validate_url_format("https://api.example.co.uk/v1") is True
+
+    # Invalid URLs
+    assert validate_url_format("example.com") is False
+    assert validate_url_format("http://invalid") is False
+    assert validate_url_format("https://example.com space") is False
+    assert validate_url_format("") is False

tests/test_main.py

@@ -1,10 +1,10 @@
 """Test cases for main module."""
 
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 import typer
-from click.core import Command, Context
+from click.core import Command
 from typer.testing import CliRunner
 
 from irrexplorer_cli.main import app, asn, prefix
@@ -69,10 +69,36 @@ def test_asn_command_without_context() -> None:
 def test_prefix_direct_no_args() -> None:
     """Test prefix function directly with no arguments."""
     with pytest.raises(typer.Exit):
-        prefix("", "", Context(Command("test")))
+        ctx = typer.Context(Command("test"))
+        ctx.obj = {"base_url": "http://example.com"}
+        prefix(ctx)
 
 
 def test_asn_direct_no_args() -> None:
     """Test asn function directly with no arguments."""
     with pytest.raises(typer.Exit):
-        asn("", "", Context(Command("test")))
+        ctx = typer.Context(Command("test"))
+        ctx.obj = {"base_url": "http://example.com"}
+        asn(ctx)
+
+
+@patch("irrexplorer_cli.main.validate_url_format")
+def test_prefix_invalid_url_format(mock_validate_url: MagicMock) -> None:
+    """Test prefix command with invalid URL format."""
+    mock_validate_url.return_value = False
+    ctx = typer.Context(Command("test"))
+    ctx.obj = {"base_url": "http://example.com"}
+    with pytest.raises(typer.Exit) as exc_info:
+        prefix(ctx, "192.0.2.0/24")
+    assert exc_info.value.exit_code == 1
+
+
+@patch("irrexplorer_cli.main.validate_url_format")
+def test_asn_invalid_url_format(mock_validate_url: MagicMock) -> None:
+    """Test ASN command with invalid URL format."""
+    mock_validate_url.return_value = False
+    ctx = typer.Context(Command("test"))
+    ctx.obj = {"base_url": "http://example.com"}
+    with pytest.raises(typer.Exit) as exc_info:
+        asn(ctx, "AS12345")
+    assert exc_info.value.exit_code == 1

tests/test_queries.py

@@ -6,9 +6,10 @@
 
 import httpx
 import pytest
+import typer
 
 from irrexplorer_cli.irrexplorer import IrrExplorer
-from irrexplorer_cli.queries import async_asn_query, process_overlaps
+from irrexplorer_cli.queries import async_asn_query, async_prefix_query, process_overlaps
 
 
 @pytest.mark.asyncio
@@ -48,3 +49,37 @@ async def test_process_overlaps_error_handling() -> None:
     with patch("irrexplorer_cli.irrexplorer.IrrExplorer.fetch_prefix_info", side_effect=httpx.HTTPError("Test error")):
         await process_overlaps(explorer, "192.0.2.0/24")
     await explorer.close()
+
+
+@pytest.mark.asyncio
+async def test_asn_query_connection_error() -> None:
+    """Test ASN query with connection error."""
+    with (
+        patch("irrexplorer_cli.irrexplorer.IrrExplorer.fetch_asn_info", side_effect=httpx.ConnectError("Test error")),
+        patch("builtins.print") as mock_print,
+        pytest.raises(typer.Exit) as exc_info,
+    ):
+        await async_asn_query("AS12345", base_url="https://example.com")
+        mock_print.assert_called_with(
+            "Error: Unable to connect to https://example.com. "
+            "Please verify the URL is correct and the service is available."
+        )
+        assert exc_info.value.exit_code == 1
+
+
+@pytest.mark.asyncio
+async def test_prefix_query_connection_error() -> None:
+    """Test prefix query with connection error."""
+    with (
+        patch(
+            "irrexplorer_cli.irrexplorer.IrrExplorer.fetch_prefix_info", side_effect=httpx.ConnectError("Test error")
+        ),
+        patch("builtins.print") as mock_print,
+        pytest.raises(typer.Exit) as exc_info,
+    ):
+        await async_prefix_query("192.0.2.0/24", base_url="https://example.com")
+        mock_print.assert_called_with(
+            "Error: Unable to connect to https://example.com. "
+            "Please verify the URL is correct and the service is available."
+        )
+        assert exc_info.value.exit_code == 1