From 4e5c5e42c8db93d1635c4ae37780e95e5e77c5c6 Mon Sep 17 00:00:00 2001 From: joon Date: Wed, 5 Jun 2024 01:12:41 -0700 Subject: [PATCH 1/4] Use black --- .flake8 | 4 +++- pyproject.toml | 12 +++++++++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/.flake8 b/.flake8 index 6deafc2..a77cb08 100644 --- a/.flake8 +++ b/.flake8 @@ -1,2 +1,4 @@ [flake8] -max-line-length = 120 +max-line-length = 80 +extend-select = B950 +extend-ignore = E203,E501,E701 diff --git a/pyproject.toml b/pyproject.toml index ea7205d..b9352c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,7 +60,9 @@ disable = true [tool.hatch.envs.default] dependencies = [ "bandit", + "black~=24.4.2", "flake8>=7.0.0", + "flake8-bugbear~=24.4.6", "freezegun", "mypy", "pytest", @@ -85,12 +87,16 @@ python = ["38", "39", "310", "311", "312"] detached = true dependencies = [ # Make sure the respective versions are synced with default! "bandit", - "flake8>=6.0.0", + "black~=24.4.2", + "flake8>=7.0.0", + "flake8-bugbear~=24.4.6", ] [tool.hatch.envs.lint.scripts] -flake = "flake8 {args:wtfis}" -security = "bandit --quiet -r {args:wtfis}" +blackcheck = "black --check wtfis tests" +flake = "flake8 wtfis tests" +security = "bandit --quiet -r wtfis" all = [ + "blackcheck", "flake", "security", ] From 1f8e3f7a8c415a25e7d457c7d34bddffc342bf55 Mon Sep 17 00:00:00 2001 From: joon Date: Mon, 10 Jun 2024 02:50:14 -0700 Subject: [PATCH 2/4] Apply black to code --- pyproject.toml | 2 +- tests/conftest.py | 17 +- tests/test_cli.py | 309 ++++++++++++++++---------- tests/test_clients.py | 16 +- tests/test_handlers.py | 175 +++++++++++---- tests/test_models_vt.py | 4 +- tests/test_ui_domain_view.py | 403 +++++++++++++++++++++------------- tests/test_ui_ip_view.py | 202 +++++++++++------ tests/test_utils.py | 5 +- wtfis/clients/abuseipdb.py | 1 + wtfis/clients/base.py | 9 +- wtfis/clients/greynoise.py | 1 + wtfis/clients/ip2whois.py | 8 +- wtfis/clients/ipwhois.py | 1 + wtfis/clients/passivetotal.py | 6 +- wtfis/clients/shodan.py | 1 + wtfis/clients/types.py | 5 +- wtfis/clients/urlhaus.py | 3 +- wtfis/clients/virustotal.py | 13 +- wtfis/exceptions.py | 2 + wtfis/handlers/base.py | 11 +- wtfis/handlers/domain.py | 39 +++- wtfis/handlers/ip.py | 21 +- wtfis/main.py | 76 ++++--- wtfis/models/base.py | 6 +- wtfis/models/ip2whois.py | 15 +- wtfis/models/ipwhois.py | 1 + wtfis/models/types.py | 13 +- wtfis/models/urlhaus.py | 9 +- wtfis/models/virustotal.py | 28 +-- wtfis/ui/base.py | 294 +++++++++++++++---------- wtfis/ui/view.py | 89 +++++--- wtfis/utils.py | 17 +- 33 files changed, 1149 insertions(+), 653 deletions(-) create mode 100644 wtfis/exceptions.py diff --git a/pyproject.toml b/pyproject.toml index b9352c8..54eb576 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,7 +93,7 @@ dependencies = [ # Make sure the respective versions are synced with default! ] [tool.hatch.envs.lint.scripts] blackcheck = "black --check wtfis tests" -flake = "flake8 wtfis tests" +flake = "flake8 wtfis" security = "bandit --quiet -r wtfis" all = [ "blackcheck", diff --git a/tests/conftest.py b/tests/conftest.py index d7110b8..902ed8b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,7 +13,8 @@ class TestTheme: - """ Expected theme values for the tests """ + """Expected theme values for the tests""" + panel_title = "bold yellow" heading_h1 = "bold bright_green on dark_green" heading_h2 = "bold yellow" @@ -53,32 +54,32 @@ def open_test_data(fname: str) -> str: def abuseipdb_get_ip(ip, pool) -> AbuseIpDb: - """ Mock replacement for AbuseIpDbClient()._get_ip() """ + """Mock replacement for AbuseIpDbClient()._get_ip()""" return AbuseIpDb.model_validate(pool[ip]) def greynoise_get(ip, pool) -> GreynoiseIp: - """ Mock replacement for GreynoiseClient().get_ip() """ + """Mock replacement for GreynoiseClient().get_ip()""" return GreynoiseIp.model_validate(pool[ip]) def ipwhois_get(ip, pool) -> IpWhois: - """ Mock replacement for IpWhoisClient().get_ipwhois() """ + """Mock replacement for IpWhoisClient().get_ipwhois()""" return IpWhois.model_validate(pool[ip]) def shodan_get_ip(ip, pool) -> ShodanIp: - """ Mock replacement for ShodanClient().get_ip() """ + """Mock replacement for ShodanClient().get_ip()""" return ShodanIp.model_validate(pool[ip]) def urlhaus_get_host(entity, pool) -> UrlHaus: - """ Mock replacement for UrlHausClient()._get_host() """ + """Mock replacement for UrlHausClient()._get_host()""" return UrlHaus.model_validate(pool[entity]) def timestamp_text(ts) -> Optional[RenderableType]: - """ Standard timestamp formatting """ + """Standard timestamp formatting""" theme = TestTheme() return Text( ts, @@ -86,7 +87,7 @@ def timestamp_text(ts) -> Optional[RenderableType]: Span(10, 11, theme.timestamp_t), Span(11, 19, theme.timestamp_time), Span(19, 20, theme.timestamp_z), - ] + ], ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 5b3cd5c..90e6b5a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -6,8 +6,14 @@ import pytest from dotenv import load_dotenv from rich.console import Console -from rich.progress import (BarColumn, Progress, SpinnerColumn, - TaskProgressColumn, TextColumn, TimeElapsedColumn) +from rich.progress import ( + BarColumn, + Progress, + SpinnerColumn, + TaskProgressColumn, + TextColumn, + TimeElapsedColumn, +) from wtfis.clients.abuseipdb import AbuseIpDbClient from wtfis.clients.greynoise import GreynoiseClient @@ -17,10 +23,16 @@ from wtfis.clients.shodan import ShodanClient from wtfis.clients.urlhaus import UrlHausClient from wtfis.clients.virustotal import VTClient +from wtfis.exceptions import WtfisException from wtfis.handlers.domain import DomainHandler from wtfis.handlers.ip import IpAddressHandler -from wtfis.main import (generate_entity_handler, generate_view, main, - parse_args, parse_env) +from wtfis.main import ( + generate_entity_handler, + generate_view, + main, + parse_args, + parse_env, +) from wtfis.models.virustotal import Domain, IpAddress from wtfis.ui.view import DomainView, IpAddressView @@ -129,10 +141,13 @@ def fake_load_dotenv_ip2whois(tmp_path): class TestArgs: def test_basic(self): - with patch("sys.argv", [ - "main", - "www.example.com", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + ], + ): args = parse_args() assert args.entity == "www.example.com" assert args.max_resolutions == 3 @@ -143,142 +158,189 @@ def test_basic(self): assert args.use_urlhaus is False def test_display(self): - with patch("sys.argv", [ - "main", - "www.example.com", - "-n", - "-1", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-n", + "-1", + ], + ): args = parse_args() assert args.no_color is True assert args.one_column is True def test_max_resolutions_ok(self): - with patch("sys.argv", [ - "main", - "www.example.com", - "-m", - "8", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-m", + "8", + ], + ): args = parse_args() assert args.max_resolutions == 8 def test_max_resolutions_error_1(self, capsys): with pytest.raises(SystemExit) as e: - with patch("sys.argv", [ - "main", - "www.example.com", - "-m", - "11", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-m", + "11", + ], + ): parse_args() capture = capsys.readouterr() - assert capture.err == "usage: main [-h]\nmain: error: Maximum --max-resolutions value is 10\n" + assert ( + capture.err + == "usage: main [-h]\nmain: error: Maximum --max-resolutions value is 10\n" + ) assert e.type == SystemExit assert e.value.code == 2 def test_max_resolutions_error_2(self, capsys): with pytest.raises(SystemExit) as e: - with patch("sys.argv", [ - "main", - "1.1.1.1", - "-m", - "5", - ]): + with patch( + "sys.argv", + [ + "main", + "1.1.1.1", + "-m", + "5", + ], + ): parse_args() capture = capsys.readouterr() - assert capture.err == "usage: main [-h]\nmain: error: --max-resolutions is not applicable to IPs\n" + assert capture.err == ( + "usage: main [-h]\nmain: error: --max-resolutions is not " + "applicable to IPs\n" + ) assert e.type == SystemExit assert e.value.code == 2 def test_shodan_ok(self): os.environ["SHODAN_API_KEY"] = "foo" - with patch("sys.argv", [ - "main", - "www.example.com", - "-s", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-s", + ], + ): args = parse_args() assert args.use_shodan is True del os.environ["SHODAN_API_KEY"] def test_shodan_error(self, capsys): with pytest.raises(SystemExit) as e: - with patch("sys.argv", [ - "main", - "www.example.com", - "-s", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-s", + ], + ): parse_args() capture = capsys.readouterr() - assert capture.err == "usage: main [-h]\nmain: error: SHODAN_API_KEY is not set\n" + assert ( + capture.err == "usage: main [-h]\nmain: error: SHODAN_API_KEY is not set\n" + ) assert e.type == SystemExit assert e.value.code == 2 def test_greynoise_ok(self): os.environ["GREYNOISE_API_KEY"] = "foo" - with patch("sys.argv", [ - "main", - "www.example.com", - "-g", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-g", + ], + ): args = parse_args() assert args.use_greynoise is True del os.environ["GREYNOISE_API_KEY"] def test_greynoise_error(self, capsys): with pytest.raises(SystemExit) as e: - with patch("sys.argv", [ - "main", - "www.example.com", - "-g", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-g", + ], + ): parse_args() capture = capsys.readouterr() - assert capture.err == "usage: main [-h]\nmain: error: GREYNOISE_API_KEY is not set\n" + assert ( + capture.err + == "usage: main [-h]\nmain: error: GREYNOISE_API_KEY is not set\n" + ) assert e.type == SystemExit assert e.value.code == 2 def test_urlhaus(self): - with patch("sys.argv", [ - "main", - "www.example.com", - "-u", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-u", + ], + ): args = parse_args() assert args.use_urlhaus is True def test_abuseipdb_ok(self): os.environ["ABUSEIPDB_API_KEY"] = "foo" - with patch("sys.argv", [ - "main", - "1.1.1.1", - "-a", - ]): + with patch( + "sys.argv", + [ + "main", + "1.1.1.1", + "-a", + ], + ): args = parse_args() assert args.use_abuseipdb is True del os.environ["ABUSEIPDB_API_KEY"] def test_abuseipdb_error(self, capsys): with pytest.raises(SystemExit) as e: - with patch("sys.argv", [ - "main", - "1.1.1.1", - "-a", - ]): + with patch( + "sys.argv", + [ + "main", + "1.1.1.1", + "-a", + ], + ): parse_args() capture = capsys.readouterr() - assert capture.err == "usage: main [-h]\nmain: error: ABUSEIPDB_API_KEY is not set\n" + assert ( + capture.err + == "usage: main [-h]\nmain: error: ABUSEIPDB_API_KEY is not set\n" + ) assert e.type == SystemExit assert e.value.code == 2 @@ -314,7 +376,8 @@ def test_error(self, mock_exists, capsys): assert capture.err == ( "Error: Environment variable VT_API_KEY not set\n" - f"Env file {Path().home() / '.env.wtfis'} was not found either. Did you forget?\n" + f"Env file {Path().home() / '.env.wtfis'} was not found either. " + "Did you forget?\n" ) assert e.type == SystemExit assert e.value.code == 1 @@ -323,10 +386,13 @@ def test_error(self, mock_exists, capsys): class TestDefaults: def test_defaults_1(self, fake_load_dotenv_2): with patch("wtfis.main.load_dotenv", fake_load_dotenv_2): - with patch("sys.argv", [ - "main", - "www.example.com", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + ], + ): parse_env() args = parse_args() assert args.entity == "www.example.com" @@ -340,11 +406,14 @@ def test_defaults_1(self, fake_load_dotenv_2): def test_defaults_2(self, fake_load_dotenv_2): with patch("wtfis.main.load_dotenv", fake_load_dotenv_2): - with patch("sys.argv", [ - "main", - "www.example.com", - "-s", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + "-s", + ], + ): parse_env() args = parse_args() assert args.entity == "www.example.com" @@ -358,10 +427,13 @@ def test_defaults_2(self, fake_load_dotenv_2): def test_defaults_3(self, fake_load_dotenv_3): with patch("wtfis.main.load_dotenv", fake_load_dotenv_3): - with patch("sys.argv", [ - "main", - "www.example.com", - ]): + with patch( + "sys.argv", + [ + "main", + "www.example.com", + ], + ): parse_env() args = parse_args() assert args.entity == "www.example.com" @@ -375,11 +447,14 @@ def test_defaults_3(self, fake_load_dotenv_3): def test_defaults_4(self, fake_load_dotenv_4): with patch("wtfis.main.load_dotenv", fake_load_dotenv_4): - with patch("sys.argv", [ - "main", - "1.1.1.1", - "-u", - ]): + with patch( + "sys.argv", + [ + "main", + "1.1.1.1", + "-u", + ], + ): parse_env() args = parse_args() assert args.entity == "1.1.1.1" @@ -393,14 +468,15 @@ def test_defaults_4(self, fake_load_dotenv_4): class TestGenEntityHandler: - """ Tests for the generate_entity_handler function """ + """Tests for the generate_entity_handler function""" + @patch("sys.argv", ["main", "www.example[.]com"]) def test_handler_domain_1(self, fake_load_dotenv_1): - """ Domain with default params """ + """Domain with default params""" with patch("wtfis.main.load_dotenv", fake_load_dotenv_1): parse_env() console = Console() - progress = simulate_progress(console), + progress = (simulate_progress(console),) entity = generate_entity_handler(parse_args(), console, progress) assert isinstance(entity, DomainHandler) assert entity.entity == "www.example.com" @@ -418,11 +494,11 @@ def test_handler_domain_1(self, fake_load_dotenv_1): @patch("sys.argv", ["main", "www.example[.]com", "-s", "-g", "-u", "-m", "5"]) def test_handler_domain_2(self, fake_load_dotenv_1): - """ Domain with Shodan and Greynoise and non-default max_resolutions """ + """Domain with Shodan and Greynoise and non-default max_resolutions""" with patch("wtfis.main.load_dotenv", fake_load_dotenv_1): parse_env() console = Console() - progress = simulate_progress(console), + progress = (simulate_progress(console),) entity = generate_entity_handler(parse_args(), console, progress) assert entity.max_resolutions == 5 assert isinstance(entity._geoasn, IpWhoisClient) @@ -434,33 +510,33 @@ def test_handler_domain_2(self, fake_load_dotenv_1): @patch("sys.argv", ["main", "www.example[.]com"]) def test_handler_domain_3(self, fake_load_dotenv_vt_whois): - """ Domain using default Ip2Whois for whois """ + """Domain using default Ip2Whois for whois""" with patch("wtfis.main.load_dotenv", fake_load_dotenv_vt_whois): parse_env() console = Console() - progress = simulate_progress(console), + progress = (simulate_progress(console),) entity = generate_entity_handler(parse_args(), console, progress) assert isinstance(entity._whois, VTClient) unset_env_vars() @patch("sys.argv", ["main", "www.example[.]com"]) def test_handler_domain_4(self, fake_load_dotenv_ip2whois): - """ Domain using default Ip2Whois for whois """ + """Domain using default Ip2Whois for whois""" with patch("wtfis.main.load_dotenv", fake_load_dotenv_ip2whois): parse_env() console = Console() - progress = simulate_progress(console), + progress = (simulate_progress(console),) entity = generate_entity_handler(parse_args(), console, progress) assert isinstance(entity._whois, Ip2WhoisClient) unset_env_vars() @patch("sys.argv", ["main", "1[.]1[.]1[.]1"]) def test_handler_ip_1(self, fake_load_dotenv_1): - """ IP with default params """ + """IP with default params""" with patch("wtfis.main.load_dotenv", fake_load_dotenv_1): parse_env() console = Console() - progress = simulate_progress(console), + progress = (simulate_progress(console),) entity = generate_entity_handler(parse_args(), console, progress) assert isinstance(entity, IpAddressHandler) assert entity.entity == "1.1.1.1" @@ -475,11 +551,11 @@ def test_handler_ip_1(self, fake_load_dotenv_1): @patch("sys.argv", ["main", "1[.]1[.]1[.]1", "-s", "-g", "-u", "-a"]) def test_handler_ip_2(self, fake_load_dotenv_1): - """ IP with various options """ + """IP with various options""" with patch("wtfis.main.load_dotenv", fake_load_dotenv_1): parse_env() console = Console() - progress = simulate_progress(console), + progress = (simulate_progress(console),) entity = generate_entity_handler(parse_args(), console, progress) assert isinstance(entity._geoasn, IpWhoisClient) assert isinstance(entity._whois, PTClient) @@ -491,10 +567,11 @@ def test_handler_ip_2(self, fake_load_dotenv_1): class TestGenView: - """ Tests for the generate_view function """ + """Tests for the generate_view function""" + @patch("wtfis.main.DomainView", return_value=MagicMock(spec=DomainView)) def test_view_domain_1(self, m_domain_view, test_data): - """ Domain view with default params """ + """Domain view with default params""" entity = DomainHandler( entity=MagicMock(), console=MagicMock(), @@ -507,7 +584,9 @@ def test_view_domain_1(self, m_domain_view, test_data): abuseipdb_client=MagicMock(), urlhaus_client=MagicMock(), ) - entity.vt_info = Domain.model_validate(json.loads(test_data("vt_domain_gist.json"))) + entity.vt_info = Domain.model_validate( + json.loads(test_data("vt_domain_gist.json")) + ) entity.whois = MagicMock() entity.ip_enrich = MagicMock() view = generate_view(MagicMock(), MagicMock(), entity) @@ -515,7 +594,7 @@ def test_view_domain_1(self, m_domain_view, test_data): @patch("wtfis.main.IpAddressView", return_value=MagicMock(spec=IpAddressView)) def test_view_ip_1(self, m_ip_view, test_data): - """ IP address view with default params """ + """IP address view with default params""" entity = IpAddressHandler( entity=MagicMock(), console=MagicMock(), @@ -528,15 +607,17 @@ def test_view_ip_1(self, m_ip_view, test_data): abuseipdb_client=MagicMock(), urlhaus_client=MagicMock(), ) - entity.vt_info = IpAddress.model_validate(json.loads(test_data("vt_ip_1.1.1.1.json"))) + entity.vt_info = IpAddress.model_validate( + json.loads(test_data("vt_ip_1.1.1.1.json")) + ) entity.whois = MagicMock() entity.ip_enrich = MagicMock() view = generate_view(MagicMock(), MagicMock(), entity) assert isinstance(view, IpAddressView) def test_view_error(self): - """ IP address view with default params """ - with pytest.raises(Exception): + """IP address view with default params""" + with pytest.raises(WtfisException): generate_view(MagicMock(), MagicMock(), "foobar") @@ -547,8 +628,10 @@ class TestMain: @patch("wtfis.main.get_progress") @patch("wtfis.main.generate_entity_handler", return_value=MagicMock()) @patch("wtfis.main.generate_view", return_value=MagicMock()) - def test_main_default(self, m_view, m_handler, m_progress, m_args, m_console, fake_load_dotenv_1): - """ Test all calls with default values """ + def test_main_default( + self, m_view, m_handler, m_progress, m_args, m_console, fake_load_dotenv_1 + ): + """Test all calls with default values""" m_args.return_value = parse_args() m_progress.return_value = simulate_progress(m_console()) with patch("wtfis.main.load_dotenv", fake_load_dotenv_1): diff --git a/tests/test_clients.py b/tests/test_clients.py index 806dc34..d051f5a 100644 --- a/tests/test_clients.py +++ b/tests/test_clients.py @@ -65,7 +65,9 @@ def test_init(self, abuseipdb_client): def test_enrich_ips(self, mock_requests_get, test_data, abuseipdb_client): mock_resp = MagicMock() mock_resp.status_code = 200 - mock_resp.json.return_value = json.loads(test_data("abuseipdb_1.1.1.1_raw.json")) + mock_resp.json.return_value = json.loads( + test_data("abuseipdb_1.1.1.1_raw.json") + ) mock_requests_get.return_value = mock_resp abuseipdb = abuseipdb_client.enrich_ips("thisdoesntmatter").root["1.1.1.1"] @@ -86,7 +88,9 @@ def test_init(self, ip2whois_client): def test_get_whois(self, mock_requests_get, test_data, ip2whois_client): mock_resp = MagicMock() mock_resp.status_code = 200 - mock_resp.json.return_value = json.loads(test_data("ip2whois_whois_hotmail.json")) + mock_resp.json.return_value = json.loads( + test_data("ip2whois_whois_hotmail.json") + ) mock_requests_get.return_value = mock_resp whois = ip2whois_client.get_whois("thisdoesntmatter") @@ -167,7 +171,7 @@ def test_init(self, shodan_client): @patch.object(Shodan, "host") def test_get_ip_apierror_invalid_key(self, mock_shodan_host, shodan_client): - """ Test invalid API key APIError """ + """Test invalid API key APIError""" mock_shodan_host.side_effect = APIError("Invalid API key") with pytest.raises(APIError) as e: @@ -178,7 +182,7 @@ def test_get_ip_apierror_invalid_key(self, mock_shodan_host, shodan_client): @patch.object(Shodan, "host") def test_get_ip_apierror_other(self, mock_shodan_host, shodan_client): - """ Test other invalid API key APIError """ + """Test other invalid API key APIError""" mock_shodan_host.side_effect = APIError("Some other error") with pytest.raises(APIError) as e: @@ -196,7 +200,9 @@ def test_init(self, urlhaus_client): def test_enrich_ips(self, mock_requests_post, test_data, urlhaus_client): mock_resp = MagicMock() mock_resp.status_code = 200 - mock_resp.json.return_value = json.loads(test_data("urlhaus_1.1.1.1.json"))["1.1.1.1"] + mock_resp.json.return_value = json.loads(test_data("urlhaus_1.1.1.1.json"))[ + "1.1.1.1" + ] mock_requests_post.return_value = mock_resp urlhaus = urlhaus_client.enrich_ips("thisdoesntmatter").root["1.1.1.1"] diff --git a/tests/test_handlers.py b/tests/test_handlers.py index b7a0b9a..0af4951 100644 --- a/tests/test_handlers.py +++ b/tests/test_handlers.py @@ -67,9 +67,11 @@ def test_entity_refang(self, domain_handler): assert handler.entity == "www.example.com" def test_fetch_data_1(self, domain_handler, test_data): - """ Test with max_resolutions = 3 (default) """ + """Test with max_resolutions = 3 (default)""" handler = domain_handler() - handler.resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + handler.resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) handler._fetch_vt_domain = MagicMock() handler._fetch_vt_resolutions = MagicMock() @@ -89,7 +91,7 @@ def test_fetch_data_1(self, domain_handler, test_data): handler._fetch_urlhaus.assert_called_once() def test_fetch_data_2(self, domain_handler): - """ Test with max_resolutions = 0 """ + """Test with max_resolutions = 0""" handler = domain_handler(0) handler._fetch_vt_domain = MagicMock() handler._fetch_vt_resolutions = MagicMock() @@ -129,7 +131,9 @@ def test_vt_http_error(self, mock_requests_get, domain_handler, capsys): capture = capsys.readouterr() - assert capture.err == "Error fetching data: 401 Client Error: None for url: None\n" + assert ( + capture.err == "Error fetching data: 401 Client Error: None for url: None\n" + ) assert e.type == SystemExit assert e.value.code == 1 @@ -141,8 +145,8 @@ def test_vt_http_error(self, mock_requests_get, domain_handler, capsys): @patch.object(requests.Session, "get") def test_vt_validation_error(self, mock_requests_get, domain_handler, capsys): """ - Test a pydantic data model ValidationError from the VT client. This also tests the - common_exception_handler decorator. + Test a pydantic data model ValidationError from the VT client. This also tests + the common_exception_handler decorator. """ handler = domain_handler() mock_resp = requests.models.Response() @@ -160,7 +164,8 @@ def test_vt_validation_error(self, mock_requests_get, domain_handler, capsys): assert capture.err.startswith( "Data model validation error: 1 validation error for Domain\ndata\n" - " Field required [type=missing, input_value={'intentionally': 'wrong data'}, input_type=dict]\n" + " Field required [type=missing, input_value={'intentionally': " + "'wrong data'}, input_type=dict]\n" ) assert e.type == SystemExit assert e.value.code == 1 @@ -171,9 +176,13 @@ def test_vt_validation_error(self, mock_requests_get, domain_handler, capsys): assert e.value.code == 1 @patch.object(requests.Session, "get") - def test_ipwhois_http_error(self, mock_requests_get, domain_handler, capsys, test_data): + def test_ipwhois_http_error( + self, mock_requests_get, domain_handler, capsys, test_data + ): handler = domain_handler() - handler.resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + handler.resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) mock_resp = requests.models.Response() @@ -181,11 +190,15 @@ def test_ipwhois_http_error(self, mock_requests_get, domain_handler, capsys, tes mock_requests_get.return_value = mock_resp handler._fetch_geoasn(*["1.2.3.4", "1.2.3.5"]) - assert handler.warnings[0].startswith("Could not fetch IPWhois: 500 Server Error:") + assert handler.warnings[0].startswith( + "Could not fetch IPWhois: 500 Server Error:" + ) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch IPWhois: 500 Server Error:") + assert capture.out.startswith( + "WARN: Could not fetch IPWhois: 500 Server Error:" + ) @patch.object(requests.Session, "get") def test_ipwhois_validation_error(self, mock_requests_get, domain_handler, capsys): @@ -194,7 +207,10 @@ def test_ipwhois_validation_error(self, mock_requests_get, domain_handler, capsy with patch.object(mock_resp, "json") as mock_resp_json: mock_resp.status_code = 200 - mock_resp_json.return_value = {"success": True, "intentionally": "wrong data"} + mock_resp_json.return_value = { + "success": True, + "intentionally": "wrong data", + } mock_requests_get.return_value = mock_resp with pytest.raises(SystemExit) as e: @@ -221,7 +237,9 @@ def test_whois_http_error(self, mock_requests_get, domain_handler, capsys): capture = capsys.readouterr() - assert capture.err == "Error fetching data: 401 Client Error: None for url: None\n" + assert ( + capture.err == "Error fetching data: 401 Client Error: None for url: None\n" + ) assert e.type == SystemExit assert e.value.code == 1 @@ -237,17 +255,22 @@ def test_vt_resolutions_429_error(self, mock_requests_get, domain_handler, capsy mock_requests_get.return_value = mock_resp handler._fetch_vt_resolutions() - assert handler.warnings[0].startswith("Could not fetch Virustotal resolutions: 429 Client Error:") + assert handler.warnings[0].startswith( + "Could not fetch Virustotal resolutions: 429 Client Error:" + ) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch Virustotal resolutions: 429 Client Error:") + assert capture.out.startswith( + "WARN: Could not fetch Virustotal resolutions: 429 Client Error:" + ) @patch.object(requests.Session, "get") def test_whois_429_error(self, mock_requests_get, domain_handler, capsys): """ Test fail open behavior of whois fetching when rate limited - Since _fetch_whois() is in the base class, this should also cover the IpAddressHandler use case. + Since _fetch_whois() is in the base class, this should also cover the + IpAddressHandler use case. """ handler = domain_handler() mock_resp = requests.models.Response() @@ -256,14 +279,18 @@ def test_whois_429_error(self, mock_requests_get, domain_handler, capsys): mock_requests_get.return_value = mock_resp handler._fetch_whois() - assert handler.warnings[0].startswith("Could not fetch Whois: 429 Client Error:") + assert handler.warnings[0].startswith( + "Could not fetch Whois: 429 Client Error:" + ) handler.print_warnings() capture = capsys.readouterr() assert capture.out.startswith("WARN: Could not fetch Whois: 429 Client Error:") @patch.object(requests.Session, "get") - def test_greynoise_429_error(self, mock_requests_get, domain_handler, capsys, test_data): + def test_greynoise_429_error( + self, mock_requests_get, domain_handler, capsys, test_data + ): """ Test fail open behavior of Greynoise when rate limited """ @@ -273,24 +300,33 @@ def test_greynoise_429_error(self, mock_requests_get, domain_handler, capsys, te mock_resp.status_code = 429 mock_requests_get.return_value = mock_resp - handler.resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + handler.resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) handler._fetch_greynoise(*["1.2.3.4", "1.2.3.5"]) - assert handler.warnings[0].startswith("Could not fetch Greynoise: 429 Client Error:") + assert handler.warnings[0].startswith( + "Could not fetch Greynoise: 429 Client Error:" + ) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch Greynoise: 429 Client Error:") + assert capture.out.startswith( + "WARN: Could not fetch Greynoise: 429 Client Error:" + ) @patch.object(requests.Session, "get") def test_greynoise_404_error(self, mock_requests_get, domain_handler, test_data): """ - Test fail open behavior of Greynoise when no IP found (404) and no warning message + Test fail open behavior of Greynoise when no IP found (404) and no warning + message """ handler = domain_handler(3) mock_resp = requests.models.Response() - handler.resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + handler.resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) mock_resp.status_code = 404 mock_requests_get.return_value = mock_resp @@ -300,44 +336,60 @@ def test_greynoise_404_error(self, mock_requests_get, domain_handler, test_data) assert handler.greynoise == GreynoiseIpMap.model_validate({}) @patch.object(requests.Session, "get") - def test_greynoise_403_error(self, mock_requests_get, domain_handler, test_data, capsys): + def test_greynoise_403_error( + self, mock_requests_get, domain_handler, test_data, capsys + ): """ Test exception behavior of Greynoise when not under any of the handled cases """ handler = domain_handler(3) mock_resp = requests.models.Response() - handler.resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + handler.resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) mock_resp.status_code = 403 mock_requests_get.return_value = mock_resp handler._fetch_greynoise(*["1.2.3.4", "1.2.3.5"]) - assert handler.warnings[0].startswith("Could not fetch Greynoise: 403 Client Error:") + assert handler.warnings[0].startswith( + "Could not fetch Greynoise: 403 Client Error:" + ) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch Greynoise: 403 Client Error:") + assert capture.out.startswith( + "WARN: Could not fetch Greynoise: 403 Client Error:" + ) @patch.object(requests.Session, "get") - def test_greynoise_connection_error(self, mock_requests_get, domain_handler, test_data, capsys): + def test_greynoise_connection_error( + self, mock_requests_get, domain_handler, test_data, capsys + ): """ Test exception behavior of Greynoise with non-HTTPError requests exception """ handler = domain_handler(3) mock_requests_get.side_effect = ConnectionError("Foo bar message") - handler.resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + handler.resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) handler._fetch_greynoise(*["1.2.3.4", "1.2.3.5"]) assert handler.warnings[0] == "Could not fetch Greynoise: Foo bar message" handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch Greynoise: Foo bar message") + assert capture.out.startswith( + "WARN: Could not fetch Greynoise: Foo bar message" + ) @patch.object(requests.Session, "get") - def test_greynoise_validation_error(self, mock_requests_get, domain_handler, capsys): + def test_greynoise_validation_error( + self, mock_requests_get, domain_handler, capsys + ): handler = domain_handler() mock_resp = requests.models.Response() @@ -411,7 +463,9 @@ def test_vt_http_error(self, mock_requests_get, ip_handler, capsys): capture = capsys.readouterr() - assert capture.err == "Error fetching data: 404 Client Error: None for url: None\n" + assert ( + capture.err == "Error fetching data: 404 Client Error: None for url: None\n" + ) assert e.type == SystemExit assert e.value.code == 1 @@ -433,7 +487,8 @@ def test_vt_validation_error(self, mock_requests_get, ip_handler, capsys): assert capture.err.startswith( "Data model validation error: 1 validation error for IpAddress\ndata\n" - " Field required [type=missing, input_value={'intentionally': 'wrong data'}, input_type=dict]\n" + " Field required [type=missing, input_value={'intentionally': " + "'wrong data'}, input_type=dict]\n" ) assert e.type == SystemExit assert e.value.code == 1 @@ -448,11 +503,15 @@ def test_ipwhois_http_error(self, mock_requests_get, ip_handler, capsys): mock_requests_get.return_value = mock_resp handler._fetch_geoasn("1.2.3.4") - assert handler.warnings[0].startswith("Could not fetch IPWhois: 502 Server Error:") + assert handler.warnings[0].startswith( + "Could not fetch IPWhois: 502 Server Error:" + ) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch IPWhois: 502 Server Error:") + assert capture.out.startswith( + "WARN: Could not fetch IPWhois: 502 Server Error:" + ) @patch.object(requests.Session, "get") def test_whois_http_error(self, mock_requests_get, ip_handler, capsys): @@ -467,7 +526,9 @@ def test_whois_http_error(self, mock_requests_get, ip_handler, capsys): capture = capsys.readouterr() - assert capture.err == "Error fetching data: 403 Client Error: None for url: None\n" + assert ( + capture.err == "Error fetching data: 403 Client Error: None for url: None\n" + ) assert e.type == SystemExit assert e.value.code == 1 @@ -478,7 +539,10 @@ def test_ipwhois_validation_error(self, mock_requests_get, ip_handler, capsys): with patch.object(mock_resp, "json") as mock_resp_json: mock_resp.status_code = 200 - mock_resp_json.return_value = {"success": True, "intentionally": "wrong data"} + mock_resp_json.return_value = { + "success": True, + "intentionally": "wrong data", + } mock_requests_get.return_value = mock_resp with pytest.raises(SystemExit) as e: @@ -501,7 +565,7 @@ def test_shodan_api_error(self, mock_requests_get, ip_handler, capsys): mock_resp = requests.models.Response() mock_resp.status_code = 401 - mock_resp._content = b'' + mock_resp._content = b"" mock_requests_get.return_value = mock_resp handler._fetch_shodan("1.2.3.4") @@ -527,7 +591,9 @@ def test_greynoise_http_error(self, mock_requests_get, ip_handler, capsys): capture = capsys.readouterr() - assert capture.err == "Error fetching data: 401 Client Error: None for url: None\n" + assert ( + capture.err == "Error fetching data: 401 Client Error: None for url: None\n" + ) assert e.type == SystemExit assert e.value.code == 1 @@ -543,16 +609,21 @@ def test_greynoise_429_error(self, mock_requests_get, ip_handler, capsys): mock_requests_get.return_value = mock_resp handler._fetch_greynoise("1.2.3.4") - assert handler.warnings[0].startswith("Could not fetch Greynoise: 429 Client Error:") + assert handler.warnings[0].startswith( + "Could not fetch Greynoise: 429 Client Error:" + ) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch Greynoise: 429 Client Error:") + assert capture.out.startswith( + "WARN: Could not fetch Greynoise: 429 Client Error:" + ) @patch.object(requests.Session, "get") def test_greynoise_404_error(self, mock_requests_get, ip_handler): """ - Test fail open behavior of Greynoise when no IP found (404) and no warning message + Test fail open behavior of Greynoise when no IP found (404) and no warning + message """ handler = ip_handler() mock_resp = requests.models.Response() @@ -576,11 +647,15 @@ def test_greynoise_403_error(self, mock_requests_get, ip_handler, capsys): mock_requests_get.return_value = mock_resp handler._fetch_greynoise("1.2.3.4") - assert handler.warnings[0].startswith("Could not fetch Greynoise: 403 Client Error:") + assert handler.warnings[0].startswith( + "Could not fetch Greynoise: 403 Client Error:" + ) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch Greynoise: 403 Client Error:") + assert capture.out.startswith( + "WARN: Could not fetch Greynoise: 403 Client Error:" + ) @patch.object(requests.Session, "get") def test_greynoise_connection_error(self, mock_requests_get, ip_handler, capsys): @@ -595,7 +670,9 @@ def test_greynoise_connection_error(self, mock_requests_get, ip_handler, capsys) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch Greynoise: Foo bar message") + assert capture.out.startswith( + "WARN: Could not fetch Greynoise: Foo bar message" + ) @patch.object(requests.Session, "get") def test_greynoise_validation_error(self, mock_requests_get, ip_handler, capsys): @@ -645,8 +722,12 @@ def test_abuseipdb_429_error(self, mock_requests_get, ip_handler, capsys): mock_requests_get.return_value = mock_resp handler._fetch_abuseipdb("1.2.3.4") - assert handler.warnings[0].startswith("Could not fetch AbuseIPDB: 429 Client Error:") + assert handler.warnings[0].startswith( + "Could not fetch AbuseIPDB: 429 Client Error:" + ) handler.print_warnings() capture = capsys.readouterr() - assert capture.out.startswith("WARN: Could not fetch AbuseIPDB: 429 Client Error:") + assert capture.out.startswith( + "WARN: Could not fetch AbuseIPDB: 429 Client Error:" + ) diff --git a/tests/test_models_vt.py b/tests/test_models_vt.py index c48edd5..a7bffe5 100644 --- a/tests/test_models_vt.py +++ b/tests/test_models_vt.py @@ -21,7 +21,9 @@ def test_domain_1(self, test_data): ] def test_resolutions_1(self, test_data): - res = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + res = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) assert res.meta.count == 37 assert len(res.data) == 10 diff --git a/tests/test_ui_domain_view.py b/tests/test_ui_domain_view.py index 6f15b16..03f1468 100644 --- a/tests/test_ui_domain_view.py +++ b/tests/test_ui_domain_view.py @@ -29,12 +29,16 @@ @pytest.fixture() def view01(test_data, mock_ipwhois_get): - """ gist.github.com with PT whois. Complete test of all panels. Also test print(). """ - resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + """gist.github.com with PT whois. Complete test of all panels. Also test print().""" + resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) geoasn_pool = json.loads(test_data("ipwhois_gist.json")) geoasn_client = IpWhoisClient() - geoasn_client._get_ipwhois = MagicMock(side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool)) + geoasn_client._get_ipwhois = MagicMock( + side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool) + ) geoasn_enrich = geoasn_client.enrich_ips(*resolutions.ip_list(3)) return DomainView( @@ -59,7 +63,9 @@ def view02(test_data): return DomainView( console=Console(), entity=MagicMock(), - resolutions=Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))), + resolutions=Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ), geoasn=IpWhoisMap.model_validate({}), whois=VTWhois.model_validate(json.loads(test_data("vt_whois_gist.json"))), shodan=ShodanIpMap.model_validate({}), @@ -72,7 +78,7 @@ def view02(test_data): @pytest.fixture() def view03(test_data): - """ bbc.co.uk VT whois. Whois panel test only. Test whois with no domain field. """ + """bbc.co.uk VT whois. Whois panel test only. Test whois with no domain field.""" return DomainView( console=Console(), entity=MagicMock(), @@ -107,7 +113,8 @@ def view04(test_data): @pytest.fixture() def view05(test_data): - """ tucows.com domain. Domain test only. Test domain with negative reputation and no popularity.""" + """tucows.com domain. Domain test only. Test domain with negative reputation and no + popularity.""" return DomainView( console=Console(), entity=Domain.model_validate(json.loads(test_data("vt_domain_tucows.json"))), @@ -123,7 +130,7 @@ def view05(test_data): @pytest.fixture() def view06(test_data): - """ exmple.com VT whois. Whois test only. Test empty whois_map.""" + """exmple.com VT whois. Whois test only. Test empty whois_map.""" return DomainView( console=Console(), entity=MagicMock(), @@ -139,17 +146,23 @@ def view06(test_data): @pytest.fixture() def view07(test_data, mock_ipwhois_get, mock_shodan_get_ip): - """ gist.github.com with Shodan. Only test resolution, geoasn and Shodan. """ - resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + """gist.github.com with Shodan. Only test resolution, geoasn and Shodan.""" + resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) geoasn_pool = json.loads(test_data("ipwhois_gist.json")) geoasn_client = IpWhoisClient() - geoasn_client._get_ipwhois = MagicMock(side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool)) + geoasn_client._get_ipwhois = MagicMock( + side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool) + ) geoasn_enrich = geoasn_client.enrich_ips(*resolutions.ip_list(3)) shodan_pool = json.loads(test_data("shodan_gist.json")) shodan_client = ShodanClient(MagicMock()) - shodan_client._get_ip = MagicMock(side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool)) + shodan_client._get_ip = MagicMock( + side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool) + ) shodan_enrich = shodan_client.enrich_ips(*resolutions.ip_list(3)) return DomainView( @@ -167,12 +180,16 @@ def view07(test_data, mock_ipwhois_get, mock_shodan_get_ip): @pytest.fixture() def view08(test_data, mock_shodan_get_ip): - """ www.wired.com with Shodan. Only test resolution and Shodan. """ - resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_wired.json"))) + """www.wired.com with Shodan. Only test resolution and Shodan.""" + resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_wired.json")) + ) shodan_pool = json.loads(test_data("shodan_wired.json")) shodan_client = ShodanClient(MagicMock()) - shodan_client._get_ip = MagicMock(side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool)) + shodan_client._get_ip = MagicMock( + side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool) + ) shodan_enrich = shodan_client.enrich_ips(*resolutions.ip_list(1)) return DomainView( @@ -191,22 +208,31 @@ def view08(test_data, mock_shodan_get_ip): @pytest.fixture() def view09(test_data, mock_shodan_get_ip, mock_greynoise_get, mock_abuseipdb_get): - """ one.one.one.one with Shodan, Greynoise and AbuseIPDB. Only test mentioned services. """ - resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_one.json"))) + """one.one.one.one with Shodan, Greynoise and AbuseIPDB. Only test mentioned + services.""" + resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_one.json")) + ) shodan_pool = json.loads(test_data("shodan_one.json")) shodan_client = ShodanClient(MagicMock()) - shodan_client._get_ip = MagicMock(side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool)) + shodan_client._get_ip = MagicMock( + side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool) + ) shodan_enrich = shodan_client.enrich_ips(*resolutions.ip_list(1)) greynoise_pool = json.loads(test_data("greynoise_one.json")) greynoise_client = GreynoiseClient("dummykey") - greynoise_client._get_ip = MagicMock(side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool)) + greynoise_client._get_ip = MagicMock( + side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool) + ) greynoise_enrich = greynoise_client.enrich_ips(*resolutions.ip_list(1)) abuseipdb_pool = json.loads(test_data("abuseipdb_one.json")) abuseipdb_client = AbuseIpDbClient("dummykey") - abuseipdb_client._get_ip = MagicMock(side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool)) + abuseipdb_client._get_ip = MagicMock( + side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool) + ) abuseipdb_enrich = abuseipdb_client.enrich_ips(*resolutions.ip_list(1)) return DomainView( @@ -225,7 +251,7 @@ def view09(test_data, mock_shodan_get_ip, mock_greynoise_get, mock_abuseipdb_get @pytest.fixture() def view10(test_data): - """ Dummy VT whois. Whois panel test only. Test whois with no data. """ + """Dummy VT whois. Whois panel test only. Test whois with no data.""" return DomainView( console=Console(), entity=MagicMock(), @@ -241,12 +267,16 @@ def view10(test_data): @pytest.fixture() def view11(test_data, mock_shodan_get_ip): - """ gist.github.com with Shodan. Only test Shodan. Test empty open ports. """ - resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + """gist.github.com with Shodan. Only test Shodan. Test empty open ports.""" + resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) shodan_pool = json.loads(test_data("shodan_gist_2.json")) shodan_client = ShodanClient(MagicMock()) - shodan_client._get_ip = MagicMock(side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool)) + shodan_client._get_ip = MagicMock( + side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool) + ) shodan_enrich = shodan_client.enrich_ips(*resolutions.ip_list(3)) return DomainView( @@ -264,13 +294,15 @@ def view11(test_data, mock_shodan_get_ip): @pytest.fixture() def view12(test_data): - """ Dummy IP2WHOIS whois. Whois panel test only. """ + """Dummy IP2WHOIS whois. Whois panel test only.""" return DomainView( console=Console(), entity=MagicMock(), resolutions=MagicMock(), geoasn=MagicMock(), - whois=Ip2Whois.model_validate(json.loads(test_data("ip2whois_whois_hotmail.json"))), + whois=Ip2Whois.model_validate( + json.loads(test_data("ip2whois_whois_hotmail.json")) + ), shodan=MagicMock(), greynoise=MagicMock(), abuseipdb=MagicMock(), @@ -280,7 +312,7 @@ def view12(test_data): @pytest.fixture() def view13(test_data): - """ Dummy IP2WHOIS whois. Whois panel test only. Test null registrant. """ + """Dummy IP2WHOIS whois. Whois panel test only. Test null registrant.""" return DomainView( console=Console(), entity=MagicMock(), @@ -296,17 +328,23 @@ def view13(test_data): @pytest.fixture() def view14(test_data, mock_ipwhois_get, mock_urlhaus_get): - """ Same as view01() but with Urlhaus enrichment. Test URLhaus only. """ - resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))) + """Same as view01() but with Urlhaus enrichment. Test URLhaus only.""" + resolutions = Resolutions.model_validate( + json.loads(test_data("vt_resolutions_gist.json")) + ) geoasn_pool = json.loads(test_data("ipwhois_gist.json")) geoasn_client = IpWhoisClient() - geoasn_client._get_ipwhois = MagicMock(side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool)) + geoasn_client._get_ipwhois = MagicMock( + side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool) + ) geoasn_enrich = geoasn_client.enrich_ips(*resolutions.ip_list(3)) urlhaus_pool = json.loads(test_data("urlhaus_gist.json")) urlhaus_client = UrlHausClient() - urlhaus_client._get_host = MagicMock(side_effect=lambda domain: mock_urlhaus_get(domain, urlhaus_pool)) + urlhaus_client._get_host = MagicMock( + side_effect=lambda domain: mock_urlhaus_get(domain, urlhaus_pool) + ) urlhaus_enrich = urlhaus_client.enrich_domains("gist.github.com") return DomainView( @@ -347,7 +385,9 @@ def test_domain_panel(self, view01, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/domain/gist.github.com')], + spans=[ + Span(0, 8, "link https://virustotal.com/gui/domain/gist.github.com") + ], ), "Reputation:", "Popularity:", @@ -367,10 +407,13 @@ def test_domain_panel(self, view01, theme, display_timestamp): Span(10, 13, "cyan"), Span(15, 29, "bright_cyan"), Span(31, 36, "cyan"), - ] + ], ), Text( - "advice, file sharing/storage, information technology, media sharing, social networks", + ( + "advice, file sharing/storage, information technology, " + "media sharing, social networks" + ), spans=[ Span(0, 6, "bright_white on black"), Span(6, 8, "default"), @@ -381,7 +424,7 @@ def test_domain_panel(self, view01, theme, display_timestamp): Span(54, 67, "bright_white on black"), Span(67, 69, "default"), Span(69, 84, "bright_white on black"), - ] + ], ), display_timestamp("2022-08-16T06:14:59Z"), display_timestamp("2022-08-15T22:25:30Z"), @@ -417,7 +460,11 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): assert group[1].columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/13.234.210.38')], + spans=[ + Span( + 0, 8, "link https://virustotal.com/gui/ip-address/13.234.210.38" + ) + ], ), "Resolved:", "ASN:", @@ -430,8 +477,7 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): Text("0/94 malicious", spans=[Span(0, 14, theme.info)]), display_timestamp("2022-08-06T14:56:20Z"), Text( - "16509 (Amazon Data Services India)", - spans=[Span(7, 33, theme.asn_org)] + "16509 (Amazon Data Services India)", spans=[Span(7, 33, theme.asn_org)] ), "Amazon.com, Inc.", Text( @@ -439,8 +485,8 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): spans=[ Span(6, 8, "default"), Span(19, 21, "default"), - ] - ) + ], + ), ] # Spacing @@ -461,7 +507,13 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): assert group[1].columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/192.30.255.113')], + spans=[ + Span( + 0, + 8, + "link https://virustotal.com/gui/ip-address/192.30.255.113", + ) + ], ), "Resolved:", "ASN:", @@ -473,18 +525,15 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): assert group[1].columns[1]._cells == [ Text("1/94 malicious", spans=[Span(0, 14, theme.error)]), display_timestamp("2022-06-21T18:10:54Z"), - Text( - "36459 (GitHub, Inc.)", - spans=[Span(7, 19, theme.asn_org)] - ), + Text("36459 (GitHub, Inc.)", spans=[Span(7, 19, theme.asn_org)]), "GitHub, Inc.", Text( "Seattle, Washington, United States", spans=[ Span(7, 9, "default"), Span(19, 21, "default"), - ] - ) + ], + ), ] # Spacing @@ -499,8 +548,8 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): spans=[Span(0, 14, theme.heading_h2)], ) - # Unlike the previous entries, the table is inside a group of (Table, Text) due to - # old timestamp warning + # Unlike the previous entries, the table is inside a group of (Table, Text) + # due to old timestamp warning # table = group[1].renderables[0] table = group[1] assert table.columns[0].style == theme.table_field @@ -508,7 +557,13 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/13.234.176.102')], + spans=[ + Span( + 0, + 8, + "link https://virustotal.com/gui/ip-address/13.234.176.102", + ) + ], ), "Resolved:", "ASN:", @@ -521,8 +576,7 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): Text("0/94 malicious", spans=[Span(0, 14, theme.info)]), display_timestamp("2015-08-17T07:11:53Z"), Text( - "16509 (Amazon Data Services India)", - spans=[Span(7, 33, theme.asn_org)] + "16509 (Amazon Data Services India)", spans=[Span(7, 33, theme.asn_org)] ), "Amazon.com, Inc.", Text( @@ -530,8 +584,8 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): spans=[ Span(6, 8, "default"), Span(19, 21, "default"), - ] - ) + ], + ), ] # Old timestamp warning @@ -540,7 +594,9 @@ def test_resolutions_panel(self, view01, theme, display_timestamp): # Spacing and remaining count footer assert res.renderable.renderables[6] == "" assert str(footer) == "+34 more" - assert footer.spans[0].style.startswith(f"{theme.footer} link https://virustotal.com/gui/domain/") + assert footer.spans[0].style.startswith( + f"{theme.footer} link https://virustotal.com/gui/domain/" + ) def test_whois_panel(self, view01, theme, display_timestamp): whois = view01.whois_panel() @@ -558,7 +614,16 @@ def test_whois_panel(self, view01, theme, display_timestamp): # Content assert content.renderables[0] == Text( "github.com", - spans=[Span(0, 10, 'bold yellow link https://community.riskiq.com/search/github.com/whois')] + spans=[ + Span( + 0, + 10, + ( + "bold yellow link " + "https://community.riskiq.com/search/github.com/whois" + ), + ) + ], ) table = content.renderables[1] @@ -585,10 +650,7 @@ def test_whois_panel(self, view01, theme, display_timestamp): assert table.columns[1].justify == "left" assert table.columns[1]._cells == [ "MarkMonitor Inc.", - Text( - "GitHub, Inc.", - spans=[] - ), + Text("GitHub, Inc.", spans=[]), "N/A", "abusecomplaints@markmonitor.com", "+1.5555555", @@ -598,8 +660,12 @@ def test_whois_panel(self, view01, theme, display_timestamp): "US", "00000", Text( - ("dns1.p08.nsone.net, dns2.p08.nsone.net, dns3.p08.nsone.net, dns4.p08.nsone.net, " - "ns-1283.awsdns-32.org, ns-1707.awsdns-21.co.uk, ns-421.awsdns-52.com, ns-520.awsdns-01.net"), + ( + "dns1.p08.nsone.net, dns2.p08.nsone.net, dns3.p08.nsone.net, " + "dns4.p08.nsone.net, ns-1283.awsdns-32.org, " + "ns-1707.awsdns-21.co.uk, ns-421.awsdns-52.com, " + "ns-520.awsdns-01.net" + ), spans=[ Span(0, 18, theme.nameserver_list), Span(18, 20, "default"), @@ -615,8 +681,8 @@ def test_whois_panel(self, view01, theme, display_timestamp): Span(126, 128, "default"), Span(128, 148, theme.nameserver_list), Span(148, 150, "default"), - Span(150, 170, theme.nameserver_list) - ] + Span(150, 170, theme.nameserver_list), + ], ), display_timestamp("2007-10-09T18:20:50Z"), display_timestamp("2020-09-08T09:18:27Z"), @@ -665,7 +731,11 @@ def test_resolutions_panel(self, view02, theme, display_timestamp): assert group[1].columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/13.234.210.38')], + spans=[ + Span( + 0, 8, "link https://virustotal.com/gui/ip-address/13.234.210.38" + ) + ], ), "Resolved:", ] @@ -679,7 +749,9 @@ def test_resolutions_panel(self, view02, theme, display_timestamp): # Spacing and remaining count footer assert res.renderable.renderables[2] == "" assert str(footer) == "+36 more" - assert footer.spans[0].style.startswith(f"{theme.footer} link https://virustotal.com/gui/domain/") + assert footer.spans[0].style.startswith( + f"{theme.footer} link https://virustotal.com/gui/domain/" + ) def test_whois_panel(self, view02, theme): whois = view02.whois_panel() @@ -696,8 +768,7 @@ def test_whois_panel(self, view02, theme): # Content assert content.renderables[0] == Text( - "github.com", - spans=[Span(0, 10, 'bold yellow')] + "github.com", spans=[Span(0, 10, "bold yellow")] ) table = content.renderables[1] @@ -714,7 +785,9 @@ def test_whois_panel(self, view02, theme): ] assert table.columns[1].style == theme.table_value assert table.columns[1].justify == "left" - assert [str(c) for c in table.columns[1]._cells] == [ # Just test the strings (no style) + assert [ + str(c) for c in table.columns[1]._cells + ] == [ # Just test the strings (no style) "MarkMonitor Inc.", "dns1.p08.nsone.net, dns2.p08.nsone.net, dns3.p08.nsone.net, " "dns4.p08.nsone.net, ns-1283.awsdns-32.org, ns-1707.awsdns-21.co.uk, " @@ -786,7 +859,7 @@ def test_domain_panel(self, view04, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/domain/google.com')], + spans=[Span(0, 8, "link https://virustotal.com/gui/domain/google.com")], ), "Reputation:", "Popularity:", @@ -802,11 +875,14 @@ def test_domain_panel(self, view04, theme, display_timestamp): spans=[ Span(0, 14, theme.error), Span(15, 20, "cyan"), - ] + ], ), Text("448"), Text( - "Majestic (1)\nStatvoo (1)\nAlexa (1)\nCisco Umbrella (2)\nQuantcast (1)", + ( + "Majestic (1)\nStatvoo (1)\nAlexa (1)\nCisco Umbrella (2)\n" + "Quantcast (1)" + ), spans=[ Span(0, 8, "bright_cyan"), Span(10, 11, "cyan"), @@ -818,10 +894,13 @@ def test_domain_panel(self, view04, theme, display_timestamp): Span(51, 52, "cyan"), Span(54, 63, "bright_cyan"), Span(65, 66, "cyan"), - ] + ], ), Text( - "mobile communications, portals, search engines, search engines and portals, searchengines", + ( + "mobile communications, portals, search engines, " + "search engines and portals, searchengines" + ), spans=[ Span(0, 21, "bright_white on black"), Span(21, 23, "default"), @@ -832,7 +911,7 @@ def test_domain_panel(self, view04, theme, display_timestamp): Span(48, 74, "bright_white on black"), Span(74, 76, "default"), Span(76, 89, "bright_white on black"), - ] + ], ), display_timestamp("2022-08-17T06:03:03Z"), display_timestamp("2022-08-17T00:35:19Z"), @@ -867,7 +946,7 @@ def test_domain_panel(self, view05, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/domain/tucows.com')], + spans=[Span(0, 8, "link https://virustotal.com/gui/domain/tucows.com")], ), "Reputation:", "Categories:", @@ -882,13 +961,14 @@ def test_domain_panel(self, view05, theme, display_timestamp): spans=[ Span(0, 14, theme.error), Span(15, 21, "cyan"), - ] + ], ), Text("-1"), Text( ( - "ads/analytics, dynamic dns and isp sites, hosting, information technology, " - "known infection source, mobile communications, not recommended site" + "ads/analytics, dynamic dns and isp sites, hosting, " + "information technology, known infection source, " + "mobile communications, not recommended site" ), spans=[ Span(0, 13, "bright_white on black"), @@ -904,7 +984,7 @@ def test_domain_panel(self, view05, theme, display_timestamp): Span(99, 120, "bright_white on black"), Span(120, 122, "default"), Span(122, 142, "bright_white on black"), - ] + ], ), display_timestamp("2022-08-17T05:30:23Z"), display_timestamp("2022-08-16T22:24:18Z"), @@ -961,7 +1041,11 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): assert group[1].columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/13.234.210.38')], + spans=[ + Span( + 0, 8, "link https://virustotal.com/gui/ip-address/13.234.210.38" + ) + ], ), "Resolved:", "ASN:", @@ -969,7 +1053,7 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): "Location:", Text( "Services:", - spans=[Span(0, 8, "link https://www.shodan.io/host/13.234.210.38")] + spans=[Span(0, 8, "link https://www.shodan.io/host/13.234.210.38")], ), "Tags:", ] @@ -978,14 +1062,14 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): assert group[1].columns[1]._cells == [ Text("0/94 malicious", spans=[Span(0, 14, theme.info)]), display_timestamp("2022-08-06T14:56:20Z"), - Text("16509 (Amazon Data Services India)", spans=[Span(7, 33, "bright_white")]), + Text( + "16509 (Amazon Data Services India)", + spans=[Span(7, 33, "bright_white")], + ), "Amazon.com, Inc.", Text( "Mumbai, Maharashtra, India", - spans=[ - Span(6, 8, "default"), - Span(19, 21, "default") - ], + spans=[Span(6, 8, "default"), Span(19, 21, "default")], ), Text( "22/tcp, 80/tcp, 443/tcp", @@ -998,14 +1082,9 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): Span(14, 16, "default"), Span(16, 19, theme.port), Span(19, 23, theme.transport), - ] - ), - Text( - "cloud", - spans=[ - Span(0, 5, 'bright_white on black') - ] + ], ), + Text("cloud", spans=[Span(0, 5, "bright_white on black")]), ] # Spacing @@ -1026,7 +1105,13 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): assert group[1].columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/192.30.255.113')], + spans=[ + Span( + 0, + 8, + "link https://virustotal.com/gui/ip-address/192.30.255.113", + ) + ], ), "Resolved:", "ASN:", @@ -1034,7 +1119,7 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): "Location:", Text( "Services:", - spans=[Span(0, 8, "link https://www.shodan.io/host/192.30.255.113")] + spans=[Span(0, 8, "link https://www.shodan.io/host/192.30.255.113")], ), ] assert group[1].columns[1].style == theme.table_value @@ -1046,10 +1131,7 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): "GitHub, Inc.", Text( "Seattle, Washington, United States", - spans=[ - Span(7, 9, "default"), - Span(19, 21, "default") - ], + spans=[Span(7, 9, "default"), Span(19, 21, "default")], ), Text( "22/tcp, 80/tcp, 443/tcp", @@ -1061,8 +1143,8 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): Span(10, 14, theme.transport), Span(14, 16, "default"), Span(16, 19, theme.port), - Span(19, 23, theme.transport) - ] + Span(19, 23, theme.transport), + ], ), ] @@ -1078,8 +1160,8 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): spans=[Span(0, 14, theme.heading_h2)], ) - # Unlike the previous entries, the table is inside a group of (Table, Text) due to - # old timestamp warning + # Unlike the previous entries, the table is inside a group of (Table, Text) due + # to old timestamp warning # table = group[1].renderables[0] table = group[1] assert table.columns[0].style == theme.table_field @@ -1087,7 +1169,13 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/13.234.176.102')], + spans=[ + Span( + 0, + 8, + "link https://virustotal.com/gui/ip-address/13.234.176.102", + ) + ], ), "Resolved:", "ASN:", @@ -1095,7 +1183,7 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): "Location:", Text( "Services:", - spans=[Span(0, 8, "link https://www.shodan.io/host/13.234.176.102")] + spans=[Span(0, 8, "link https://www.shodan.io/host/13.234.176.102")], ), "Tags:", ] @@ -1104,14 +1192,14 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): assert table.columns[1]._cells == [ Text("0/94 malicious", spans=[Span(0, 14, theme.info)]), display_timestamp("2015-08-17T07:11:53Z"), - Text("16509 (Amazon Data Services India)", spans=[Span(7, 33, "bright_white")]), + Text( + "16509 (Amazon Data Services India)", + spans=[Span(7, 33, "bright_white")], + ), "Amazon.com, Inc.", Text( "Mumbai, Maharashtra, India", - spans=[ - Span(6, 8, "default"), - Span(19, 21, "default") - ], + spans=[Span(6, 8, "default"), Span(19, 21, "default")], ), Text( "22/tcp, 80/tcp, 443/tcp", @@ -1124,14 +1212,9 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): Span(14, 16, "default"), Span(16, 19, theme.port), Span(19, 23, theme.transport), - ] - ), - Text( - "cloud", - spans=[ - Span(0, 5, 'bright_white on black') - ] + ], ), + Text("cloud", spans=[Span(0, 5, "bright_white on black")]), ] # Old timestamp warning @@ -1140,7 +1223,9 @@ def test_resolutions_panel(self, view07, theme, display_timestamp): # Spacing and remaining count footer assert res.renderable.renderables[6] == "" assert str(footer) == "+34 more" - assert footer.spans[0].style.startswith(f"{theme.footer} link https://virustotal.com/gui/domain/") + assert footer.spans[0].style.startswith( + f"{theme.footer} link https://virustotal.com/gui/domain/" + ) class TestView08: @@ -1172,12 +1257,18 @@ def test_resolutions_panel(self, view08, theme, display_timestamp): assert group[1].columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/199.232.34.194')], + spans=[ + Span( + 0, + 8, + "link https://virustotal.com/gui/ip-address/199.232.34.194", + ) + ], ), "Resolved:", Text( "Services:", - spans=[Span(0, 8, "link https://www.shodan.io/host/199.232.34.194")] + spans=[Span(0, 8, "link https://www.shodan.io/host/199.232.34.194")], ), "Tags:", ] @@ -1195,13 +1286,13 @@ def test_resolutions_panel(self, view08, theme, display_timestamp): Span(28, 33, theme.product), Span(35, 38, theme.port), Span(38, 42, theme.transport), - ] + ], ), Text( "cdn", spans=[ - Span(0, 3, 'bright_white on black'), - ] + Span(0, 3, "bright_white on black"), + ], ), ] @@ -1210,7 +1301,9 @@ def test_resolutions_panel(self, view08, theme, display_timestamp): # Footer assert str(footer) == "+199 more" - assert footer.spans[0].style.startswith(f"{theme.footer} link https://virustotal.com/gui/domain/") + assert footer.spans[0].style.startswith( + f"{theme.footer} link https://virustotal.com/gui/domain/" + ) class TestView09: @@ -1244,20 +1337,22 @@ def test_resolutions_panel(self, view09, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/1.0.0.1')], + spans=[ + Span(0, 8, "link https://virustotal.com/gui/ip-address/1.0.0.1") + ], ), "Resolved:", Text( "Services:", - spans=[Span(0, 8, "link https://www.shodan.io/host/1.0.0.1")] + spans=[Span(0, 8, "link https://www.shodan.io/host/1.0.0.1")], ), Text( "GreyNoise:", - spans=[Span(0, 9, "link https://viz.greynoise.io/riot/1.0.0.1")] + spans=[Span(0, 9, "link https://viz.greynoise.io/riot/1.0.0.1")], ), Text( "AbuseIPDB:", - spans=[Span(0, 9, "link https://www.abuseipdb.com/check/1.0.0.1")] + spans=[Span(0, 9, "link https://www.abuseipdb.com/check/1.0.0.1")], ), ] assert table.columns[1].style == theme.table_value @@ -1267,8 +1362,8 @@ def test_resolutions_panel(self, view09, theme, display_timestamp): display_timestamp("2020-08-01T22:07:20Z"), Text( ( - "CloudFlare (80/tcp, 8080/tcp)\nOther (53/tcp, 53/udp, 443/tcp, 2082/tcp, " - "2086/tcp, 2087/tcp, 8443/tcp)" + "CloudFlare (80/tcp, 8080/tcp)\nOther (53/tcp, 53/udp, 443/tcp, " + "2082/tcp, 2086/tcp, 2087/tcp, 8443/tcp)" ), spans=[ Span(0, 10, theme.product), @@ -1298,7 +1393,7 @@ def test_resolutions_panel(self, view09, theme, display_timestamp): Span(90, 92, "default"), Span(92, 96, theme.port), Span(96, 100, theme.transport), - ] + ], ), Text( "✓ riot ✗ noise ✓ benign", @@ -1309,7 +1404,7 @@ def test_resolutions_panel(self, view09, theme, display_timestamp): Span(10, 15, theme.tags), Span(17, 18, theme.info), Span(19, 25, theme.tags_green), - ] + ], ), Text( "100 confidence score (567 reports)", @@ -1317,7 +1412,7 @@ def test_resolutions_panel(self, view09, theme, display_timestamp): Span(0, 3, theme.error), Span(3, 20, "red"), Span(20, 34, theme.table_value), - ] + ], ), ] @@ -1329,7 +1424,9 @@ def test_resolutions_panel(self, view09, theme, display_timestamp): # Footer assert str(footer) == "+1 more" - assert footer.spans[0].style.startswith(f"{theme.footer} link https://virustotal.com/gui/domain/") + assert footer.spans[0].style.startswith( + f"{theme.footer} link https://virustotal.com/gui/domain/" + ) class TestView10: @@ -1379,7 +1476,11 @@ def test_resolutions_panel(self, view11, theme, display_timestamp): assert group[1].columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/13.234.210.38')], + spans=[ + Span( + 0, 8, "link https://virustotal.com/gui/ip-address/13.234.210.38" + ) + ], ), "Resolved:", "Tags:", @@ -1389,12 +1490,7 @@ def test_resolutions_panel(self, view11, theme, display_timestamp): assert group[1].columns[1]._cells == [ Text("0/94 malicious", spans=[Span(0, 14, theme.info)]), display_timestamp("2022-08-06T14:56:20Z"), - Text( - "cloud", - spans=[ - Span(0, 5, theme.tags) - ] - ), + Text("cloud", spans=[Span(0, 5, theme.tags)]), ] @@ -1414,8 +1510,7 @@ def test_whois_panel(self, view12, theme): # Content assert content.renderables[0] == Text( - "hotmail.com", - spans=[Span(0, 11, "bold yellow")] + "hotmail.com", spans=[Span(0, 11, "bold yellow")] ) # Table @@ -1441,7 +1536,9 @@ def test_whois_panel(self, view12, theme): ] assert table.columns[1].style == theme.table_value assert table.columns[1].justify == "left" - assert [str(c) for c in table.columns[1]._cells] == [ # Just test the strings (no style) + assert [ + str(c) for c in table.columns[1]._cells + ] == [ # Just test the strings (no style) "MarkMonitor, Inc.", "Microsoft Corporation", "Domain Administrator", @@ -1452,8 +1549,10 @@ def test_whois_panel(self, view12, theme): "WA", "US", "98052", - ("ns4-205.azure-dns.info, ns3-205.azure-dns.org, ns1-205.azure-dns.com, " - "ns2-205.azure-dns.net"), + ( + "ns4-205.azure-dns.info, ns3-205.azure-dns.org, ns1-205.azure-dns.com, " + "ns2-205.azure-dns.net" + ), "1996-03-27T05:00:00Z", "2021-02-02T17:08:19Z", "2024-03-27T07:00:00Z", @@ -1487,7 +1586,9 @@ def test_whois_panel(self, view13, theme): ] assert table.columns[1].style == theme.table_value assert table.columns[1].justify == "left" - assert [str(c) for c in table.columns[1]._cells] == [ # Just test the strings (no style) + assert [ + str(c) for c in table.columns[1]._cells + ] == [ # Just test the strings (no style) "1996-08-01T00:00:00Z", "2020-12-10T00:00:00Z", "2025-12-13T00:00:00Z", @@ -1511,7 +1612,9 @@ def test_domain_panel_urlhaus(self, view14, theme): assert table.columns[0]._cells == [ Text( "Malware URLs:", - spans=[Span(0, 12, "link https://urlhaus.abuse.ch/host/gist.github.com/")], + spans=[ + Span(0, 12, "link https://urlhaus.abuse.ch/host/gist.github.com/") + ], ), "Blocklists:", "Tags:", @@ -1533,7 +1636,7 @@ def test_domain_panel_urlhaus(self, view14, theme): Span(24, 32, theme.urlhaus_bl_name), Span(33, 39, theme.urlhaus_bl_high), Span(43, 48, theme.urlhaus_bl_name), - ] + ], ), Text( "Pikabot, TA577, foo, zip", @@ -1545,6 +1648,6 @@ def test_domain_panel_urlhaus(self, view14, theme): Span(16, 19, theme.tags), Span(19, 21, "default"), Span(21, 24, theme.tags), - ] + ], ), ] diff --git a/tests/test_ui_ip_view.py b/tests/test_ui_ip_view.py index bd74e39..36cb04d 100644 --- a/tests/test_ui_ip_view.py +++ b/tests/test_ui_ip_view.py @@ -24,33 +24,50 @@ @pytest.fixture() -def view01(test_data, mock_abuseipdb_get, mock_ipwhois_get, mock_shodan_get_ip, mock_greynoise_get, mock_urlhaus_get): - """ 1.1.1.1 with PT whois. Complete test of all panels. Also test print(). """ +def view01( + test_data, + mock_abuseipdb_get, + mock_ipwhois_get, + mock_shodan_get_ip, + mock_greynoise_get, + mock_urlhaus_get, +): + """1.1.1.1 with PT whois. Complete test of all panels. Also test print().""" ip = "1.1.1.1" geoasn_pool = json.loads(test_data("ipwhois_1.1.1.1.json")) geoasn_client = IpWhoisClient() - geoasn_client._get_ipwhois = MagicMock(side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool)) + geoasn_client._get_ipwhois = MagicMock( + side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool) + ) geoasn_enrich = geoasn_client.enrich_ips(ip) shodan_pool = json.loads(test_data("shodan_1.1.1.1.json")) shodan_client = ShodanClient(MagicMock()) - shodan_client._get_ip = MagicMock(side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool)) + shodan_client._get_ip = MagicMock( + side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool) + ) shodan_enrich = shodan_client.enrich_ips(ip) greynoise_pool = json.loads(test_data("greynoise_1.1.1.1.json")) greynoise_client = GreynoiseClient("dummykey") - greynoise_client._get_ip = MagicMock(side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool)) + greynoise_client._get_ip = MagicMock( + side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool) + ) greynoise_enrich = greynoise_client.enrich_ips(ip) abuseipdb_pool = json.loads(test_data("abuseipdb_1.1.1.1_red.json")) abuseipdb_client = AbuseIpDbClient("dummykey") - abuseipdb_client._get_ip = MagicMock(side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool)) + abuseipdb_client._get_ip = MagicMock( + side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool) + ) abuseipdb_enrich = abuseipdb_client.enrich_ips(ip) urlhaus_pool = json.loads(test_data("urlhaus_1.1.1.1.json")) urlhaus_client = UrlHausClient() - urlhaus_client._get_host = MagicMock(side_effect=lambda ip: mock_urlhaus_get(ip, urlhaus_pool)) + urlhaus_client._get_host = MagicMock( + side_effect=lambda ip: mock_urlhaus_get(ip, urlhaus_pool) + ) urlhaus_enrich = urlhaus_client.enrich_ips(ip) return IpAddressView( @@ -67,21 +84,27 @@ def view01(test_data, mock_abuseipdb_get, mock_ipwhois_get, mock_shodan_get_ip, @pytest.fixture() def view02(test_data, mock_ipwhois_get, mock_shodan_get_ip, mock_greynoise_get): - """ 1.1.1.1 with Shodan and Greynoise. Test the whole IP panel. """ + """1.1.1.1 with Shodan and Greynoise. Test the whole IP panel.""" ip = "1.1.1.1" geoasn_pool = json.loads(test_data("ipwhois_1.1.1.1.json")) geoasn_client = IpWhoisClient() - geoasn_client._get_ipwhois = MagicMock(side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool)) + geoasn_client._get_ipwhois = MagicMock( + side_effect=lambda ip: mock_ipwhois_get(ip, geoasn_pool) + ) geoasn_enrich = geoasn_client.enrich_ips(ip) shodan_pool = json.loads(test_data("shodan_1.1.1.1.json")) shodan_client = ShodanClient(MagicMock()) - shodan_client._get_ip = MagicMock(side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool)) + shodan_client._get_ip = MagicMock( + side_effect=lambda ip: mock_shodan_get_ip(ip, shodan_pool) + ) shodan_enrich = shodan_client.enrich_ips(ip) greynoise_pool = json.loads(test_data("greynoise_1.1.1.1.json")) greynoise_client = GreynoiseClient("dummykey") - greynoise_client._get_ip = MagicMock(side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool)) + greynoise_client._get_ip = MagicMock( + side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool) + ) greynoise_enrich = greynoise_client.enrich_ips(ip) return IpAddressView( @@ -98,7 +121,7 @@ def view02(test_data, mock_ipwhois_get, mock_shodan_get_ip, mock_greynoise_get): @pytest.fixture() def view03(test_data): - """ 1.1.1.1 VT whois. Whois panel test only.""" + """1.1.1.1 VT whois. Whois panel test only.""" return IpAddressView( console=Console(), entity=MagicMock(), @@ -119,7 +142,9 @@ def view04(test_data): """ return IpAddressView( console=Console(), - entity=IpAddress.model_validate(json.loads(test_data("vt_ip_142.251.220.110.json"))), + entity=IpAddress.model_validate( + json.loads(test_data("vt_ip_142.251.220.110.json")) + ), geoasn=IpWhoisMap.model_validate({}), whois=MagicMock(), shodan=MagicMock(), @@ -131,11 +156,13 @@ def view04(test_data): @pytest.fixture() def view05(test_data, mock_greynoise_get): - """ 1.1.1.1 with alt Greynoise results. Test Greynoise only. """ + """1.1.1.1 with alt Greynoise results. Test Greynoise only.""" ip = "1.1.1.1" greynoise_pool = json.loads(test_data("greynoise_1.1.1.1_malicious.json")) greynoise_client = GreynoiseClient("dummykey") - greynoise_client._get_ip = MagicMock(side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool)) + greynoise_client._get_ip = MagicMock( + side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool) + ) greynoise_enrich = greynoise_client.enrich_ips(ip) return IpAddressView( @@ -152,11 +179,14 @@ def view05(test_data, mock_greynoise_get): @pytest.fixture() def view06(test_data, mock_greynoise_get): - """ 1.1.1.1 with another alt Greynoise result (unknown class). Test Greynoise only. """ + """1.1.1.1 with another alt Greynoise result (unknown class). + Test Greynoise only.""" ip = "1.1.1.1" greynoise_pool = json.loads(test_data("greynoise_1.1.1.1_unknown.json")) greynoise_client = GreynoiseClient("dummykey") - greynoise_client._get_ip = MagicMock(side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool)) + greynoise_client._get_ip = MagicMock( + side_effect=lambda ip: mock_greynoise_get(ip, greynoise_pool) + ) greynoise_enrich = greynoise_client.enrich_ips(ip) return IpAddressView( @@ -173,11 +203,13 @@ def view06(test_data, mock_greynoise_get): @pytest.fixture() def view07(test_data, mock_abuseipdb_get): - """ 1.1.1.1 with green AbuseIPDB score. Test AbuseIPDB only. """ + """1.1.1.1 with green AbuseIPDB score. Test AbuseIPDB only.""" ip = "1.1.1.1" abuseipdb_pool = json.loads(test_data("abuseipdb_1.1.1.1_green.json")) abuseipdb_client = AbuseIpDbClient("dummykey") - abuseipdb_client._get_ip = MagicMock(side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool)) + abuseipdb_client._get_ip = MagicMock( + side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool) + ) abuseipdb_enrich = abuseipdb_client.enrich_ips(ip) return IpAddressView( @@ -194,11 +226,13 @@ def view07(test_data, mock_abuseipdb_get): @pytest.fixture() def view08(test_data, mock_abuseipdb_get): - """ 1.1.1.1 with yellow AbuseIPDB score. Test AbuseIPDB only. """ + """1.1.1.1 with yellow AbuseIPDB score. Test AbuseIPDB only.""" ip = "1.1.1.1" abuseipdb_pool = json.loads(test_data("abuseipdb_1.1.1.1_yellow.json")) abuseipdb_client = AbuseIpDbClient("dummykey") - abuseipdb_client._get_ip = MagicMock(side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool)) + abuseipdb_client._get_ip = MagicMock( + side_effect=lambda ip: mock_abuseipdb_get(ip, abuseipdb_pool) + ) abuseipdb_enrich = abuseipdb_client.enrich_ips(ip) return IpAddressView( @@ -248,7 +282,9 @@ def test_ip_panel(self, view01, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/1.1.1.1')], + spans=[ + Span(0, 8, "link https://virustotal.com/gui/ip-address/1.1.1.1") + ], ), "Reputation:", "Updated:", @@ -257,9 +293,12 @@ def test_ip_panel(self, view01, theme, display_timestamp): assert table.columns[1].justify == "left" assert table.columns[1]._cells == [ Text( - "4/94 malicious\nCMC Threat Intelligence, Comodo Valkyrie Verdict, CRDF, Blueliv", + ( + "4/94 malicious\nCMC Threat Intelligence, Comodo Valkyrie Verdict, " + "CRDF, Blueliv" + ), spans=[ - Span(0, 14, theme.error), + Span(0, 14, theme.error), Span(15, 38, theme.vendor_list), Span(38, 40, "default"), Span(40, 63, theme.vendor_list), @@ -267,7 +306,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): Span(65, 69, theme.vendor_list), Span(69, 71, "default"), Span(71, 78, theme.vendor_list), - ] + ], ), Text("134"), display_timestamp("2022-09-03T06:47:04Z"), @@ -304,7 +343,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): spans=[ Span(6, 8, "default"), Span(23, 25, "default"), - ] + ], ), ] @@ -324,7 +363,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Services:", - spans=[Span(0, 8, 'link https://www.shodan.io/host/1.1.1.1')], + spans=[Span(0, 8, "link https://www.shodan.io/host/1.1.1.1")], ), "Last Scan:", ] @@ -333,9 +372,11 @@ def test_ip_panel(self, view01, theme, display_timestamp): assert table.columns[1]._cells == [ Text( ( - "Cisco router tftpd (69/udp)\nCloudFlare (80/tcp, 8080/tcp, 8880/tcp)\n" - "DrayTek Vigor Router (443/tcp)\nOther (53/tcp, 53/udp, 161/udp, " - "2082/tcp, 2083/tcp, 2086/tcp, 2087/tcp, 8443/tcp)" + "Cisco router tftpd (69/udp)\n" + "CloudFlare (80/tcp, 8080/tcp, 8880/tcp)\n" + "DrayTek Vigor Router (443/tcp)\n" + "Other (53/tcp, 53/udp, 161/udp, 2082/tcp, 2083/tcp, 2086/tcp, " + "2087/tcp, 8443/tcp)" ), spans=[ Span(0, 18, theme.product), @@ -377,7 +418,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): Span(169, 171, "default"), Span(171, 175, theme.port), Span(175, 179, theme.transport), - ] + ], ), display_timestamp("2022-09-04T01:03:56Z"), ] @@ -398,7 +439,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Malware URLs:", - spans=[Span(0, 12, 'link https://urlhaus.abuse.ch/host/1.1.1.1/')] + spans=[Span(0, 12, "link https://urlhaus.abuse.ch/host/1.1.1.1/")], ), "Blocklists:", "Tags:", @@ -412,7 +453,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): spans=[ Span(0, 9, theme.error), Span(9, 20, theme.table_value), - ] + ], ), Text( "not listed in spamhaus\nnot listed in surbl", @@ -421,7 +462,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): Span(14, 22, theme.urlhaus_bl_name), Span(23, 33, theme.urlhaus_bl_low), Span(37, 42, theme.urlhaus_bl_name), - ] + ], ), Text( "elf, mirai", @@ -429,7 +470,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): Span(0, 3, theme.tags), Span(3, 5, "default"), Span(5, 10, theme.tags), - ] + ], ), ] @@ -449,12 +490,12 @@ def test_ip_panel(self, view01, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "GreyNoise:", - spans=[Span(0, 9, "link https://viz.greynoise.io/riot/1.1.1.1")] + spans=[Span(0, 9, "link https://viz.greynoise.io/riot/1.1.1.1")], ), Text( "AbuseIPDB:", - spans=[Span(0, 9, "link https://www.abuseipdb.com/check/1.1.1.1")] - ) + spans=[Span(0, 9, "link https://www.abuseipdb.com/check/1.1.1.1")], + ), ] assert table.columns[1].style == theme.table_value assert table.columns[1].justify == "left" @@ -468,7 +509,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): Span(10, 15, theme.tags), Span(17, 18, theme.info), Span(19, 25, theme.tags_green), - ] + ], ), Text( "100 confidence score (567 reports)", @@ -476,7 +517,7 @@ def test_ip_panel(self, view01, theme, display_timestamp): Span(0, 3, theme.error), Span(3, 20, "red"), Span(20, 34, theme.table_value), - ] + ], ), ] @@ -496,7 +537,16 @@ def test_whois_panel(self, view01, theme): # Heading assert content.renderables[0] == Text( "1.1.1.0", - spans=[Span(0, 7, f"{theme.heading_h2} link https://community.riskiq.com/search/1.1.1.0/whois")] + spans=[ + Span( + 0, + 7, + ( + f"{theme.heading_h2} link " + "https://community.riskiq.com/search/1.1.1.0/whois" + ), + ) + ], ) # Table @@ -570,7 +620,9 @@ def test_ip_panel(self, view02, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/1.1.1.1')], + spans=[ + Span(0, 8, "link https://virustotal.com/gui/ip-address/1.1.1.1") + ], ), "Reputation:", "Updated:", @@ -579,7 +631,10 @@ def test_ip_panel(self, view02, theme, display_timestamp): assert table.columns[1].justify == "left" assert table.columns[1]._cells == [ Text( - "4/94 malicious\nCMC Threat Intelligence, Comodo Valkyrie Verdict, CRDF, Blueliv", + ( + "4/94 malicious\nCMC Threat Intelligence, Comodo Valkyrie Verdict, " + "CRDF, Blueliv" + ), spans=[ Span(0, 14, theme.error), Span(15, 38, theme.vendor_list), @@ -589,7 +644,7 @@ def test_ip_panel(self, view02, theme, display_timestamp): Span(65, 69, theme.vendor_list), Span(69, 71, "default"), Span(71, 78, theme.vendor_list), - ] + ], ), Text("134"), display_timestamp("2022-09-03T06:47:04Z"), @@ -623,10 +678,7 @@ def test_ip_panel(self, view02, theme, display_timestamp): "Cloudflare, Inc.", Text( "Sydney, New South Wales, Australia", - spans=[ - Span(6, 8, "default"), - Span(23, 25, "default") - ], + spans=[Span(6, 8, "default"), Span(23, 25, "default")], ), ] @@ -646,7 +698,7 @@ def test_ip_panel(self, view02, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Services:", - spans=[Span(0, 8, 'link https://www.shodan.io/host/1.1.1.1')], + spans=[Span(0, 8, "link https://www.shodan.io/host/1.1.1.1")], ), "Last Scan:", ] @@ -655,7 +707,8 @@ def test_ip_panel(self, view02, theme, display_timestamp): assert table.columns[1]._cells == [ Text( ( - "Cisco router tftpd (69/udp)\nCloudFlare (80/tcp, 8080/tcp, 8880/tcp)\n" + "Cisco router tftpd (69/udp)\n" + "CloudFlare (80/tcp, 8080/tcp, 8880/tcp)\n" "DrayTek Vigor Router (443/tcp)\nOther (53/tcp, 53/udp, 161/udp, " "2082/tcp, 2083/tcp, 2086/tcp, 2087/tcp, 8443/tcp)" ), @@ -699,7 +752,7 @@ def test_ip_panel(self, view02, theme, display_timestamp): Span(169, 171, "default"), Span(171, 175, theme.port), Span(175, 179, theme.transport), - ] + ], ), display_timestamp("2022-09-04T01:03:56Z"), ] @@ -720,7 +773,7 @@ def test_ip_panel(self, view02, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "GreyNoise:", - spans=[Span(0, 9, "link https://viz.greynoise.io/riot/1.1.1.1")] + spans=[Span(0, 9, "link https://viz.greynoise.io/riot/1.1.1.1")], ), ] assert table.columns[1].style == theme.table_value @@ -735,7 +788,7 @@ def test_ip_panel(self, view02, theme, display_timestamp): Span(10, 15, theme.tags), Span(17, 18, theme.info), Span(19, 25, theme.tags_green), - ] + ], ), ] @@ -757,7 +810,7 @@ def test_whois_panel(self, view03, theme): # Content assert content.renderables[0] == Text( "one.one", - spans=[Span(0, 7, 'bold yellow')], + spans=[Span(0, 7, "bold yellow")], ) table = content.renderables[1] @@ -787,9 +840,9 @@ def test_whois_panel(self, view03, theme): "One.com A/S", "REDACTED FOR PRIVACY", ( - "Please query the RDDS service of the Registrar of Record identified in this " - "output for information on how to contact the Registrant, Admin, or Tech " - "contact of the queried domain name." + "Please query the RDDS service of the Registrar of Record identified " + "in this output for information on how to contact the Registrant, " + "Admin, or Tech contact of the queried domain name." ), "REDACTED FOR PRIVACY", "REDACTED FOR PRIVACY", @@ -797,15 +850,16 @@ def test_whois_panel(self, view03, theme): "dk", "REDACTED FOR PRIVACY", ( - "* a response from the service that a domain name is 'available', does not " - "guarantee that is able to be registered,, * we may restrict, suspend or " - "terminate your access to the service at any time, and, * the copying, " - "compilation, repackaging, dissemination or other use of the information " - "provided by the service is not permitted, without our express written " - "consent., this information has been prepared and published in order to " - "represent administrative and technical management of the tld., we may " - "discontinue or amend any part or the whole of these terms of service from " - "time to time at our absolute discretion." + "* a response from the service that a domain name is 'available', does " + "not guarantee that is able to be registered,, * we may restrict, " + "suspend or terminate your access to the service at any time, and, " + "* the copying, compilation, repackaging, dissemination or other use " + "of the information provided by the service is not permitted, without " + "our express written consent., this information has been prepared and " + "published in order to represent administrative and technical " + "management of the tld., we may discontinue or amend any part or the " + "whole of these terms of service from time to time at our absolute " + "discretion." ), "signedDelegation", "2015-05-20T12:15:44Z", @@ -838,7 +892,13 @@ def test_ip_panel(self, view04, theme, display_timestamp): assert table.columns[0]._cells == [ Text( "Analysis:", - spans=[Span(0, 8, 'link https://virustotal.com/gui/ip-address/142.251.220.110')], + spans=[ + Span( + 0, + 8, + "link https://virustotal.com/gui/ip-address/142.251.220.110", + ) + ], ), "Reputation:", "Updated:", @@ -869,7 +929,7 @@ def test_ip_panel_greynoise_only(self, view05, theme): Span(10, 15, theme.tags), Span(17, 18, theme.error), Span(19, 28, theme.tags_red), - ] + ], ) @@ -889,7 +949,7 @@ def test_ip_panel_greynoise_only(self, view06, theme): Span(8, 9, theme.warn), Span(10, 15, theme.tags), Span(19, 26, theme.tags), - ] + ], ) @@ -906,7 +966,7 @@ def test_abuseipdb_green(self, view07, theme): spans=[ Span(0, 1, theme.info), Span(1, 18, "green"), - ] + ], ) def test_abuseipdb_yellow(self, view08, theme): @@ -922,5 +982,5 @@ def test_abuseipdb_yellow(self, view08, theme): Span(0, 2, theme.warn), Span(2, 19, "yellow"), Span(19, 33, theme.table_value), - ] + ], ) diff --git a/tests/test_utils.py b/tests/test_utils.py index 9550209..8ba7b12 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -18,7 +18,10 @@ def test_Timestamp_1(self): assert Timestamp(1660428690).timestamp == "2022-08-13T22:11:30Z" def test_Timestamp_2(self): - assert Timestamp("2017-04-26T15:43:49.000-07:00").timestamp == "2017-04-26T22:43:49Z" + assert ( + Timestamp("2017-04-26T15:43:49.000-07:00").timestamp + == "2017-04-26T22:43:49Z" + ) def test_Timestamp_3(self): assert Timestamp(None).timestamp is None diff --git a/wtfis/clients/abuseipdb.py b/wtfis/clients/abuseipdb.py index 4fcce77..c6c0ad1 100644 --- a/wtfis/clients/abuseipdb.py +++ b/wtfis/clients/abuseipdb.py @@ -8,6 +8,7 @@ class AbuseIpDbClient(BaseRequestsClient, BaseIpEnricherClient): """ AbuseIPDB client """ + baseurl = "https://api.abuseipdb.com/api/v2" def __init__(self, api_key: str) -> None: diff --git a/wtfis/clients/base.py b/wtfis/clients/base.py index 640a97a..5b9c720 100644 --- a/wtfis/clients/base.py +++ b/wtfis/clients/base.py @@ -18,6 +18,7 @@ class BaseClient(abc.ABC): Base client All clients should at least inherit from this class """ + @property @abc.abstractmethod def name(self) -> str: # pragma: no coverage @@ -28,6 +29,7 @@ class BaseRequestsClient(BaseClient): """ Client that uses the requests library """ + baseurl: Union[AbstractAttribute, str] = AbstractAttribute() def __init__(self) -> None: @@ -60,6 +62,7 @@ class BaseWhoisClient(abc.ABC): """ Client used for whois lookups """ + @abc.abstractmethod def get_whois(self, entity: str) -> WhoisBase: # pragma: no coverage return NotImplemented @@ -69,8 +72,11 @@ class BaseDomainEnricherClient(abc.ABC): """ Client used for domain/FQDN enrichments """ + @abc.abstractmethod - def enrich_domains(self, *domains: str) -> DomainEnrichmentType: # pragma: no coverage + def enrich_domains( + self, *domains: str + ) -> DomainEnrichmentType: # pragma: no coverage return NotImplemented @@ -78,6 +84,7 @@ class BaseIpEnricherClient(abc.ABC): """ Client used for IP enrichments """ + @abc.abstractmethod def enrich_ips(self, *ips: str) -> IpEnrichmentType: # pragma: no coverage return NotImplemented diff --git a/wtfis/clients/greynoise.py b/wtfis/clients/greynoise.py index a3a0398..0cfc96a 100644 --- a/wtfis/clients/greynoise.py +++ b/wtfis/clients/greynoise.py @@ -9,6 +9,7 @@ class GreynoiseClient(BaseRequestsClient, BaseIpEnricherClient): """ Greynoise client """ + baseurl = "https://api.greynoise.io/v3/community" def __init__(self, api_key: str) -> None: diff --git a/wtfis/clients/ip2whois.py b/wtfis/clients/ip2whois.py index 176b265..eb3b1fa 100644 --- a/wtfis/clients/ip2whois.py +++ b/wtfis/clients/ip2whois.py @@ -9,6 +9,7 @@ class Ip2WhoisClient(BaseRequestsClient, BaseWhoisClient): """ IP2WHOIS client """ + baseurl = "https://api.ip2whois.com/v2" def __init__(self, api_key: str) -> None: @@ -29,10 +30,9 @@ def get_whois(self, domain: str) -> Whois: try: return Whois.model_validate(self._get("/", params)) except HTTPError as e: - if ( - e.response.status_code == 404 or - (e.response.status_code == 400 and - e.response.json().get("error", {})["error_code"] == 10007) + if e.response.status_code == 404 or ( + e.response.status_code == 400 + and e.response.json().get("error", {})["error_code"] == 10007 ): return Whois.model_validate({}) raise diff --git a/wtfis/clients/ipwhois.py b/wtfis/clients/ipwhois.py index 7a13115..42ab8db 100644 --- a/wtfis/clients/ipwhois.py +++ b/wtfis/clients/ipwhois.py @@ -9,6 +9,7 @@ class IpWhoisClient(BaseRequestsClient, BaseIpEnricherClient): """ IPWhois client """ + baseurl = "https://ipwho.is" @property diff --git a/wtfis/clients/passivetotal.py b/wtfis/clients/passivetotal.py index 1ab4ad7..3bb650f 100644 --- a/wtfis/clients/passivetotal.py +++ b/wtfis/clients/passivetotal.py @@ -7,6 +7,7 @@ class PTClient(BaseRequestsClient, BaseWhoisClient): """ Passivetotal client """ + baseurl = "https://api.riskiq.net/pt/v2" def __init__(self, api_user: str, api_key: str) -> None: @@ -18,10 +19,7 @@ def name(self) -> str: return "Passivetotal" def _query(self, path: str, query: str) -> dict: - return self._get( - path, - params={"query": query} - ) + return self._get(path, params={"query": query}) def get_passive_dns(self, domain: str) -> dict: return self._query("/dns/passive", refang(domain)) diff --git a/wtfis/clients/shodan.py b/wtfis/clients/shodan.py index 2573c37..3e42379 100644 --- a/wtfis/clients/shodan.py +++ b/wtfis/clients/shodan.py @@ -9,6 +9,7 @@ class ShodanClient(BaseClient, BaseIpEnricherClient): """ Shodan client """ + def __init__(self, api_key: str) -> None: self.s = Shodan(api_key) diff --git a/wtfis/clients/types.py b/wtfis/clients/types.py index 9a72a3d..1226417 100644 --- a/wtfis/clients/types.py +++ b/wtfis/clients/types.py @@ -1,6 +1,7 @@ """ Type aliases """ + from typing import Union from wtfis.clients.ip2whois import Ip2WhoisClient @@ -10,9 +11,7 @@ # IP geolocation and ASN client types -IpGeoAsnClientType = Union[ - IpWhoisClient, -] +IpGeoAsnClientType = Union[IpWhoisClient,] # IP whois client types IpWhoisClientType = Union[ diff --git a/wtfis/clients/urlhaus.py b/wtfis/clients/urlhaus.py index cfd27b1..0593096 100644 --- a/wtfis/clients/urlhaus.py +++ b/wtfis/clients/urlhaus.py @@ -10,6 +10,7 @@ class UrlHausClient(BaseRequestsClient, BaseDomainEnricherClient, BaseIpEnricher """ URLhaus client """ + baseurl = "https://urlhaus-api.abuse.ch/v1" @property @@ -20,7 +21,7 @@ def _get_host(self, host: str) -> UrlHaus: return UrlHaus.model_validate(self._post("/host", {"host": host})) def _enrich(self, *entities: str) -> UrlHausMap: - """ Method is the same whether input is a domain or IP """ + """Method is the same whether input is a domain or IP""" urlhaus_map = {} for entity in entities: data = self._get_host(entity) diff --git a/wtfis/clients/virustotal.py b/wtfis/clients/virustotal.py index 3879afb..4e32471 100644 --- a/wtfis/clients/virustotal.py +++ b/wtfis/clients/virustotal.py @@ -12,14 +12,17 @@ class VTClient(BaseRequestsClient, BaseWhoisClient): """ Virustotal client """ + baseurl = "https://www.virustotal.com/api/v3" def __init__(self, api_key: str) -> None: super().__init__() - self.s.headers.update({ - "x-apikey": api_key, - "Accept": "application/json", - }) + self.s.headers.update( + { + "x-apikey": api_key, + "Accept": "application/json", + } + ) @property def name(self) -> str: @@ -41,7 +44,7 @@ def get_ip_whois(self, ip: str) -> Whois: return Whois.model_validate(self._get(f"/ip_addresses/{ip}/historical_whois")) def get_whois(self, entity: str) -> Whois: - """ Generalized for domain and IP """ + """Generalized for domain and IP""" if is_ip(entity): return self.get_ip_whois(entity) else: diff --git a/wtfis/exceptions.py b/wtfis/exceptions.py new file mode 100644 index 0000000..2e01d68 --- /dev/null +++ b/wtfis/exceptions.py @@ -0,0 +1,2 @@ +class WtfisException(Exception): + """Custom wtfis exception""" diff --git a/wtfis/handlers/base.py b/wtfis/handlers/base.py index f79765e..5fd7674 100644 --- a/wtfis/handlers/base.py +++ b/wtfis/handlers/base.py @@ -32,7 +32,8 @@ def common_exception_handler(func: Callable) -> Callable: - """ Decorator for handling common fetch errors """ + """Decorator for handling common fetch errors""" + def inner(*args, **kwargs) -> None: progress: Progress = args[0].progress # args[0] is the method's self input try: @@ -43,11 +44,13 @@ def inner(*args, **kwargs) -> None: except ValidationError as e: progress.stop() error_and_exit(f"Data model validation error: {e}") + return inner def failopen_exception_handler(client_attr_name: str) -> Callable: - """ Decorator for handling calls that can fail open """ + """Decorator for handling calls that can fail open""" + def inner(func): def wrapper(*args, **kwargs) -> None: client = getattr(args[0], client_attr_name) # Client obj who made the call @@ -57,7 +60,9 @@ def wrapper(*args, **kwargs) -> None: except (APIError, RequestException) as e: # Add warning warnings.append(f"Could not fetch {client.name}: {e}") + return wrapper + return inner @@ -103,7 +108,7 @@ def __init__( @abc.abstractmethod def fetch_data(self) -> None: - """ Main method that controls what get fetched """ + """Main method that controls what get fetched""" return NotImplemented # type: ignore # pragma: no coverage @common_exception_handler diff --git a/wtfis/handlers/domain.py b/wtfis/handlers/domain.py index 11431ed..a39e126 100644 --- a/wtfis/handlers/domain.py +++ b/wtfis/handlers/domain.py @@ -1,6 +1,7 @@ """ Logic handler for domain and hostname inputs """ + from typing import Optional from requests.exceptions import HTTPError @@ -38,8 +39,18 @@ def __init__( urlhaus_client: Optional[UrlHausClient], max_resolutions: int = 0, ): - super().__init__(entity, console, progress, vt_client, ip_geoasn_client, whois_client, - shodan_client, greynoise_client, abuseipdb_client, urlhaus_client) + super().__init__( + entity, + console, + progress, + vt_client, + ip_geoasn_client, + whois_client, + shodan_client, + greynoise_client, + abuseipdb_client, + urlhaus_client, + ) # Extended attributes self.max_resolutions = max_resolutions @@ -76,36 +87,48 @@ def fetch_data(self): self.progress.update(task_v, completed=100) if self.resolutions and self.resolutions.data: - task_g = self.progress.add_task(f"Fetching IP location and ASN from {self._geoasn.name}") + task_g = self.progress.add_task( + f"Fetching IP location and ASN from {self._geoasn.name}" + ) self.progress.update(task_g, advance=50) self._fetch_geoasn(*self.resolutions.ip_list(self.max_resolutions)) self.progress.update(task_g, completed=100) if self._shodan: - task_s = self.progress.add_task(f"Fetching IP data from {self._shodan.name}") + task_s = self.progress.add_task( + f"Fetching IP data from {self._shodan.name}" + ) self.progress.update(task_s, advance=50) self._fetch_shodan(*self.resolutions.ip_list(self.max_resolutions)) self.progress.update(task_s, completed=100) if self._greynoise: - task_g = self.progress.add_task(f"Fetching IP data from {self._greynoise.name}") + task_g = self.progress.add_task( + f"Fetching IP data from {self._greynoise.name}" + ) self.progress.update(task_g, advance=50) self._fetch_greynoise(*self.resolutions.ip_list(self.max_resolutions)) self.progress.update(task_g, completed=100) if self._abuseipdb: - task_g = self.progress.add_task(f"Fetching IP data from {self._abuseipdb.name}") + task_g = self.progress.add_task( + f"Fetching IP data from {self._abuseipdb.name}" + ) self.progress.update(task_g, advance=50) self._fetch_abuseipdb(*self.resolutions.ip_list(self.max_resolutions)) self.progress.update(task_g, completed=100) if self._urlhaus: - task_u = self.progress.add_task(f"Fetching domain data from {self._urlhaus.name}") + task_u = self.progress.add_task( + f"Fetching domain data from {self._urlhaus.name}" + ) self.progress.update(task_u, advance=50) self._fetch_urlhaus() self.progress.update(task_u, completed=100) - task_w = self.progress.add_task(f"Fetching domain whois from {self._whois.name}") + task_w = self.progress.add_task( + f"Fetching domain whois from {self._whois.name}" + ) self.progress.update(task_w, advance=50) self._fetch_whois() self.progress.update(task_w, completed=100) diff --git a/wtfis/handlers/ip.py b/wtfis/handlers/ip.py index 118948f..2c915ae 100644 --- a/wtfis/handlers/ip.py +++ b/wtfis/handlers/ip.py @@ -1,6 +1,7 @@ """ Logic handler for IP address inputs """ + from wtfis.handlers.base import ( BaseHandler, common_exception_handler, @@ -25,31 +26,41 @@ def fetch_data(self): self._fetch_vt_ip_address() self.progress.update(task_v, completed=100) - task_g = self.progress.add_task(f"Fetching IP location and ASN from {self._geoasn.name}") + task_g = self.progress.add_task( + f"Fetching IP location and ASN from {self._geoasn.name}" + ) self.progress.update(task_g, advance=50) self._fetch_geoasn(self.entity) self.progress.update(task_g, completed=100) if self._shodan: - task_s = self.progress.add_task(f"Fetching IP data from {self._shodan.name}") + task_s = self.progress.add_task( + f"Fetching IP data from {self._shodan.name}" + ) self.progress.update(task_s, advance=50) self._fetch_shodan(self.entity) self.progress.update(task_s, completed=100) if self._urlhaus: - task_u = self.progress.add_task(f"Fetching IP data from {self._urlhaus.name}") + task_u = self.progress.add_task( + f"Fetching IP data from {self._urlhaus.name}" + ) self.progress.update(task_u, advance=50) self._fetch_urlhaus() self.progress.update(task_u, completed=100) if self._greynoise: - task_g = self.progress.add_task(f"Fetching IP data from {self._greynoise.name}") + task_g = self.progress.add_task( + f"Fetching IP data from {self._greynoise.name}" + ) self.progress.update(task_g, advance=50) self._fetch_greynoise(self.entity) self.progress.update(task_g, completed=100) if self._abuseipdb: - task_a = self.progress.add_task(f"Fetching IP data from {self._abuseipdb.name}") + task_a = self.progress.add_task( + f"Fetching IP data from {self._abuseipdb.name}" + ) self.progress.update(task_a, advance=50) self._fetch_abuseipdb(self.entity) self.progress.update(task_a, completed=100) diff --git a/wtfis/main.py b/wtfis/main.py index 51377cc..fdf421c 100644 --- a/wtfis/main.py +++ b/wtfis/main.py @@ -15,6 +15,7 @@ from wtfis.clients.shodan import ShodanClient from wtfis.clients.urlhaus import UrlHausClient from wtfis.clients.virustotal import VTClient +from wtfis.exceptions import WtfisException from wtfis.handlers.base import BaseHandler from wtfis.handlers.domain import DomainHandler from wtfis.handlers.ip import IpAddressHandler @@ -33,13 +34,14 @@ def parse_env() -> None: load_dotenv(DEFAULT_ENV_FILE) # Exit if required environment variables don't exist - for envvar in ( - "VT_API_KEY", - ): + for envvar in ("VT_API_KEY",): if not os.environ.get(envvar): error = f"Error: Environment variable {envvar} not set" if not DEFAULT_ENV_FILE.exists(): - error = error + f"\nEnv file {DEFAULT_ENV_FILE} was not found either. Did you forget?" + error += ( + f"\nEnv file {DEFAULT_ENV_FILE} was not found either. " + "Did you forget?" + ) error_and_exit(error) @@ -49,22 +51,43 @@ def parse_args() -> Namespace: parser = argparse.ArgumentParser() parser.add_argument("entity", help="Hostname, domain or IP") parser.add_argument( - "-m", "--max-resolutions", metavar="N", - help=f"Maximum number of resolutions to show (default: {DEFAULT_MAX_RESOLUTIONS})", + "-m", + "--max-resolutions", + metavar="N", + help=( + "Maximum number of resolutions to show " + f"(default: {DEFAULT_MAX_RESOLUTIONS})" + ), type=int, - default=DEFAULT_MAX_RESOLUTIONS + default=DEFAULT_MAX_RESOLUTIONS, ) - parser.add_argument("-s", "--use-shodan", help="Use Shodan to enrich IPs", action="store_true") - parser.add_argument("-g", "--use-greynoise", help="Enable Greynoise for IPs", action="store_true") - parser.add_argument("-a", "--use-abuseipdb", help="Enable AbuseIPDB for IPs", action="store_true") - parser.add_argument("-u", "--use-urlhaus", help="Enable URLhaus for IPs and domains", action="store_true") - parser.add_argument("-n", "--no-color", help="Show output without colors", action="store_true") - parser.add_argument("-1", "--one-column", help="Display results in one column", action="store_true") parser.add_argument( - "-V", "--version", + "-s", "--use-shodan", help="Use Shodan to enrich IPs", action="store_true" + ) + parser.add_argument( + "-g", "--use-greynoise", help="Enable Greynoise for IPs", action="store_true" + ) + parser.add_argument( + "-a", "--use-abuseipdb", help="Enable AbuseIPDB for IPs", action="store_true" + ) + parser.add_argument( + "-u", + "--use-urlhaus", + help="Enable URLhaus for IPs and domains", + action="store_true", + ) + parser.add_argument( + "-n", "--no-color", help="Show output without colors", action="store_true" + ) + parser.add_argument( + "-1", "--one-column", help="Display results in one column", action="store_true" + ) + parser.add_argument( + "-V", + "--version", help="Print version number", action="version", - version=get_version() + version=get_version(), ) parsed = parser.parse_args() @@ -117,8 +140,8 @@ def generate_entity_handler( # 2. IP2Whois (Domain only) # 2. Virustotal (fallback) if os.environ.get("PT_API_USER") and os.environ.get("PT_API_KEY"): - whois_client: Union[PTClient, Ip2WhoisClient, VTClient] = ( - PTClient(os.environ["PT_API_USER"], os.environ["PT_API_KEY"]) + whois_client: Union[PTClient, Ip2WhoisClient, VTClient] = PTClient( + os.environ["PT_API_USER"], os.environ["PT_API_KEY"] ) elif os.environ.get("IP2WHOIS_API_KEY") and not is_ip(args.entity): whois_client = Ip2WhoisClient(os.environ["IP2WHOIS_API_KEY"]) @@ -126,30 +149,21 @@ def generate_entity_handler( whois_client = vt_client shodan_client = ( - ShodanClient(os.environ["SHODAN_API_KEY"]) - if args.use_shodan - else None + ShodanClient(os.environ["SHODAN_API_KEY"]) if args.use_shodan else None ) # Greynoise client (optional) greynoise_client = ( - GreynoiseClient(os.environ["GREYNOISE_API_KEY"]) - if args.use_greynoise - else None + GreynoiseClient(os.environ["GREYNOISE_API_KEY"]) if args.use_greynoise else None ) # AbuseIPDB client (optional) abuseipdb_client = ( - AbuseIpDbClient(os.environ["ABUSEIPDB_API_KEY"]) - if args.use_abuseipdb else None + AbuseIpDbClient(os.environ["ABUSEIPDB_API_KEY"]) if args.use_abuseipdb else None ) # URLhaus client (optional) - urlhaus_client = ( - UrlHausClient() - if args.use_urlhaus - else None - ) + urlhaus_client = UrlHausClient() if args.use_urlhaus else None # Domain / FQDN handler if not is_ip(args.entity): @@ -215,7 +229,7 @@ def generate_view( entity.urlhaus, ) else: - raise Exception("Unsupported entity!") + raise WtfisException("Unsupported entity!") return view diff --git a/wtfis/models/base.py b/wtfis/models/base.py index 1d36424..c89bce0 100644 --- a/wtfis/models/base.py +++ b/wtfis/models/base.py @@ -17,7 +17,8 @@ class WhoisBase(BaseModel, abc.ABC): - """ Use to normalize WHOIS fields from different sources """ + """Use to normalize WHOIS fields from different sources""" + source: Optional[str] = None domain: Optional[str] = None registrar: Optional[str] = None @@ -38,7 +39,8 @@ class WhoisBase(BaseModel, abc.ABC): class IpGeoAsnBase(BaseModel, abc.ABC): - """ Use to normalize IP geolocation and ASN fields """ + """Use to normalize IP geolocation and ASN fields""" + model_config = ConfigDict(coerce_numbers_to_str=True) ip: str diff --git a/wtfis/models/ip2whois.py b/wtfis/models/ip2whois.py index 7a25fe9..81e1f9a 100644 --- a/wtfis/models/ip2whois.py +++ b/wtfis/models/ip2whois.py @@ -27,13 +27,20 @@ class Whois(WhoisBase): @model_validator(mode="before") @classmethod def extract_registrant(cls, v): - """ Surface registrant fields to root level """ + """Surface registrant fields to root level""" registrant = v.pop("registrant", {}) if not registrant: return v for field in [ - "organization", "name", "email", "phone", - "street_address", "city", "region", "country", "zip_code", + "organization", + "name", + "email", + "phone", + "street_address", + "city", + "region", + "country", + "zip_code", ]: v[field] = registrant.get(field) return v @@ -41,5 +48,5 @@ def extract_registrant(cls, v): @field_validator("registrar", mode="before") @classmethod def transform_registrar(cls, v): - """ Convert registrar from dict to simply registrar.name """ + """Convert registrar from dict to simply registrar.name""" return v.get("name") if v else v diff --git a/wtfis/models/ipwhois.py b/wtfis/models/ipwhois.py index bab1d88..9954e83 100644 --- a/wtfis/models/ipwhois.py +++ b/wtfis/models/ipwhois.py @@ -2,6 +2,7 @@ ipwhois datamodels API doc: https://ipwhois.io/documentation """ + from __future__ import annotations from typing import Dict diff --git a/wtfis/models/types.py b/wtfis/models/types.py index 73d2241..bb3538a 100644 --- a/wtfis/models/types.py +++ b/wtfis/models/types.py @@ -1,6 +1,7 @@ """ Type aliases """ + from typing import Union from wtfis.models.abuseipdb import AbuseIpDbMap @@ -20,15 +21,9 @@ ] # Domain/FQDN enrichment map types -DomainEnrichmentType = Union[ - UrlHausMap, -] +DomainEnrichmentType = Union[UrlHausMap,] # IP geolocation and ASN types -IpGeoAsnType = Union[ - IpWhois, -] +IpGeoAsnType = Union[IpWhois,] -IpGeoAsnMapType = Union[ - IpWhoisMap, -] +IpGeoAsnMapType = Union[IpWhoisMap,] diff --git a/wtfis/models/urlhaus.py b/wtfis/models/urlhaus.py index 9f27ece..20ebcda 100644 --- a/wtfis/models/urlhaus.py +++ b/wtfis/models/urlhaus.py @@ -24,14 +24,14 @@ class Url(BaseModel): @field_validator("larted", mode="before") @classmethod def convert_larted(cls, v): - """ Cast larted to bool """ + """Cast larted to bool""" mapping = {"true": True, "false": False} return mapping[v.lower()] @field_validator("tags", mode="before") @classmethod def handle_none_tags(cls, v): - """ Turn NoneType tags into an empty list """ + """Turn NoneType tags into an empty list""" return [] if v is None else v @@ -51,7 +51,7 @@ class UrlHaus(BaseModel): @field_validator("url_count", mode="before") @classmethod def convert_url_count(cls, v): - """ Cast url_count to int """ + """Cast url_count to int""" return int(v) if v else None @property @@ -59,7 +59,8 @@ def online_url_count(self) -> int: if not self._online_url_count: self._online_url_count = ( len([u for u in self.urls if u.url_status == "online"]) - if self.urls else 0 + if self.urls + else 0 ) return self._online_url_count diff --git a/wtfis/models/virustotal.py b/wtfis/models/virustotal.py index d85fe8f..e66c7f3 100644 --- a/wtfis/models/virustotal.py +++ b/wtfis/models/virustotal.py @@ -159,30 +159,30 @@ def get_latest_whois_record_and_transform(cls, v): # Normalized fields with multiple possible sources fields_w_multiple_possible_sources = { "date_changed": ( - transformed.pop("last_updated", None) or - transformed.pop("Updated Date", None) or - transformed.pop("Last updated", None) + transformed.pop("last_updated", None) + or transformed.pop("Updated Date", None) + or transformed.pop("Last updated", None) ), "date_created": ( - transformed.pop("Creation Date", None) or - transformed.pop("Registered on", None) + transformed.pop("Creation Date", None) + or transformed.pop("Registered on", None) ), "date_expires": ( - transformed.pop("Expiry Date", None) or - transformed.pop("Registry Expiry Date", None) or - transformed.pop("Expiry date", None) + transformed.pop("Expiry Date", None) + or transformed.pop("Registry Expiry Date", None) + or transformed.pop("Expiry date", None) ), "name": ( - transformed.pop("registrant_name", None) or - transformed.pop("Registrant Name", None) + transformed.pop("registrant_name", None) + or transformed.pop("Registrant Name", None) ), "country": ( - transformed.pop("registrant_country", None) or - transformed.pop("Registrant Country", None) + transformed.pop("registrant_country", None) + or transformed.pop("Registrant Country", None) ), "registrar": ( - transformed.pop("registrar_name", None) or - transformed.pop("Registrar", None) + transformed.pop("registrar_name", None) + or transformed.pop("Registrar", None) ), } diff --git a/wtfis/ui/base.py b/wtfis/ui/base.py index 4457cb2..3dbb5c7 100644 --- a/wtfis/ui/base.py +++ b/wtfis/ui/base.py @@ -10,6 +10,7 @@ from rich.table import Table from rich.text import Text from typing import Any, Generator, List, Optional, Tuple, Union +from wtfis.exceptions import WtfisException from wtfis.models.abuseipdb import AbuseIpDb, AbuseIpDbMap from wtfis.models.base import WhoisBase from wtfis.models.greynoise import GreynoiseIp, GreynoiseIpMap @@ -28,6 +29,7 @@ class BaseView(abc.ABC): """ Handles the look of the output """ + vt_gui_baseurl_domain = "https://virustotal.com/gui/domain" vt_gui_baseurl_ip = "https://virustotal.com/gui/ip-address" pt_gui_baseurl = "https://community.riskiq.com/search" @@ -56,25 +58,32 @@ def __init__( def _vendors_who_flagged_malicious(self) -> List[str]: vendors = [] - for key, result in self.entity.data.attributes.last_analysis_results.root.items(): + for ( + key, + result, + ) in self.entity.data.attributes.last_analysis_results.root.items(): if result.category == "malicious": vendors.append(key) return vendors - def _gen_heading_text(self, heading: str, hyperlink: Optional[str] = None, type: Optional[str] = "h1") -> Text: - """ Heading text - Generates 2 types: - "h1": Style is applied across the entire line - "h2": Style is applied to the text only + def _gen_heading_text( + self, heading: str, hyperlink: Optional[str] = None, type: Optional[str] = "h1" + ) -> Text: + """Heading text + Generates 2 types: + "h1": Style is applied across the entire line + "h2": Style is applied to the text only """ link_style = f" link {hyperlink}" if hyperlink else "" if type == "h1": - return Text(heading, style=f"{self.theme.heading_h1}{link_style}", justify="center") + return Text( + heading, style=f"{self.theme.heading_h1}{link_style}", justify="center" + ) elif type == "h2": text = Text(justify="center") return text.append(heading, style=f"{self.theme.heading_h2}{link_style}") else: # pragma: no cover - raise Exception(f"Invalid heading type \"{type}\"") + raise WtfisException(f'Invalid heading type "{type}"') def _gen_linked_field_name(self, name: str, hyperlink: str) -> Text: text = Text(style=self.theme.table_field) @@ -82,12 +91,16 @@ def _gen_linked_field_name(self, name: str, hyperlink: str) -> Text: text.append(":") return text - def _gen_table(self, *params: Tuple[Union[Text, str], Union[RenderableType, None]]) -> Union[Table, str]: - """ Each param should be a tuple of (field, value) """ + def _gen_table( + self, *params: Tuple[Union[Text, str], Union[RenderableType, None]] + ) -> Union[Table, str]: + """Each param should be a tuple of (field, value)""" # Set up table grid = Table.grid(expand=False, padding=(0, 1)) - grid.add_column(style=self.theme.table_field) # Field - grid.add_column(style=self.theme.table_value, max_width=38, overflow="fold") # Value + grid.add_column(style=self.theme.table_field) # Field + grid.add_column( + style=self.theme.table_value, max_width=38, overflow="fold" + ) # Value # Populate rows valid_rows = 0 @@ -107,8 +120,10 @@ def _gen_group(self, content: List[RenderableType]) -> Generator: yield item @staticmethod - def _gen_section(body: RenderableType, heading: Optional[Text] = None) -> RenderableType: - """ A section is a subset of a panel, with its own title and content """ + def _gen_section( + body: RenderableType, heading: Optional[Text] = None + ) -> RenderableType: + """A section is a subset of a panel, with its own title and content""" return Group(heading, body) if heading else body def _gen_panel( @@ -122,15 +137,19 @@ def _gen_panel( return Panel(renderable, expand=False) def _gen_vt_analysis_stats( - self, - stats: LastAnalysisStats, - vendors: Optional[List[str]] = None + self, stats: LastAnalysisStats, vendors: Optional[List[str]] = None ) -> Text: # Custom style stats_style = self.theme.error if stats.malicious >= 1 else self.theme.info # Total count - total = stats.harmless + stats.malicious + stats.suspicious + stats.timeout + stats.undetected + total = ( + stats.harmless + + stats.malicious + + stats.suspicious + + stats.timeout + + stats.undetected + ) # Text text = Text() @@ -185,10 +204,7 @@ def ports_stylized(ports: list) -> Generator: grouped = ip.group_ports_by_product() # Return a simple port list if no identified ports - if ( - len(list(grouped.keys())) == 1 and - list(grouped.keys())[0] == "Other" - ): + if len(list(grouped.keys())) == 1 and list(grouped.keys())[0] == "Other": return smart_join(*ports_stylized(grouped["Other"])) # Return grouped display of there are identified ports @@ -218,41 +234,43 @@ def _gen_greynoise_tuple(self, ip: GreynoiseIp) -> Tuple[Text, Text]: text = Text() # RIOT - riot_icon = (Text("✓", style=true_style) - if ip.riot is True - else Text("✗", style=false_style)) - (text - .append(riot_icon) - .append(" ") - .append(Text("riot", style=text_style)) - .append(" ")) + riot_icon = ( + Text("✓", style=true_style) + if ip.riot is True + else Text("✗", style=false_style) + ) + ( + text.append(riot_icon) + .append(" ") + .append(Text("riot", style=text_style)) + .append(" ") + ) # Noise - noise_icon = (Text("✓", style=true_style) - if ip.noise is True - else Text("✗", style=false_style)) - (text - .append(noise_icon) - .append(" ") - .append(Text("noise", style=text_style))) + noise_icon = ( + Text("✓", style=true_style) + if ip.noise is True + else Text("✗", style=false_style) + ) + (text.append(noise_icon).append(" ").append(Text("noise", style=text_style))) # Classification if ip.classification: text.append(" ") if ip.classification == "benign": - (text - .append(Text("✓", style=self.theme.info)) - .append(" ") - .append(Text("benign", style=self.theme.tags_green))) + ( + text.append(Text("✓", style=self.theme.info)) + .append(" ") + .append(Text("benign", style=self.theme.tags_green)) + ) elif ip.classification == "malicious": - (text - .append(Text("!", style=self.theme.error)) - .append(" ") - .append(Text("malicious", style=self.theme.tags_red))) + ( + text.append(Text("!", style=self.theme.error)) + .append(" ") + .append(Text("malicious", style=self.theme.tags_red)) + ) else: - (text - .append("? ") - .append(Text(ip.classification, style=text_style))) + (text.append("? ").append(Text(ip.classification, style=text_style))) return title, text @@ -261,7 +279,9 @@ def _gen_abuseipdb_tuple(self, ip: AbuseIpDb) -> Tuple[Text, Text]: # # Title # - title = self._gen_linked_field_name("AbuseIPDB", hyperlink=f"https://www.abuseipdb.com/check/{ip.ip_address}") + title = self._gen_linked_field_name( + "AbuseIPDB", hyperlink=f"https://www.abuseipdb.com/check/{ip.ip_address}" + ) # # Content @@ -275,9 +295,11 @@ def _gen_abuseipdb_tuple(self, ip: AbuseIpDb) -> Tuple[Text, Text]: style = self.theme.error text = Text() - (text - .append(Text(str(ip.abuse_confidence_score), style=style)) - .append(" confidence score", style=style.replace("bold ", ""))) + ( + text.append(Text(str(ip.abuse_confidence_score), style=style)).append( + " confidence score", style=style.replace("bold ", "") + ) + ) if ip.abuse_confidence_score > 0: text.append(f" ({ip.total_reports} reports)", style=self.theme.table_value) @@ -293,10 +315,11 @@ def _gen_asn_text( return None text = Text() - (text - .append(f"{asn.replace('AS', '')} (") - .append(str(org), style=self.theme.asn_org) - .append(")")) + ( + text.append(f"{asn.replace('AS', '')} (") + .append(str(org), style=self.theme.asn_org) + .append(")") + ) return text def _get_geoasn_enrichment(self, ip: str) -> Optional[IpGeoAsnType]: @@ -315,18 +338,20 @@ def _get_urlhaus_enrichment(self, entity: str) -> Optional[UrlHaus]: return self.urlhaus.root[entity] if entity in self.urlhaus.root.keys() else None def _gen_vt_section(self) -> RenderableType: - """ Virustotal section. Applies to both domain and IP views """ + """Virustotal section. Applies to both domain and IP views""" attributes = self.entity.data.attributes - baseurl = self.vt_gui_baseurl_ip if is_ip(self.entity.data.id_) else self.vt_gui_baseurl_domain + baseurl = ( + self.vt_gui_baseurl_ip + if is_ip(self.entity.data.id_) + else self.vt_gui_baseurl_domain + ) # Analysis (IP and domain) analysis = self._gen_vt_analysis_stats( - attributes.last_analysis_stats, - self._vendors_who_flagged_malicious() + attributes.last_analysis_stats, self._vendors_who_flagged_malicious() ) analysis_field = self._gen_linked_field_name( - "Analysis", - hyperlink=f"{baseurl}/{self.entity.data.id_}" + "Analysis", hyperlink=f"{baseurl}/{self.entity.data.id_}" ) # Reputation (IP and domain) @@ -346,29 +371,29 @@ def _gen_vt_section(self) -> RenderableType: # Categories (Domain only) if hasattr(attributes, "categories"): data += [ - ("Categories:", - smart_join(*attributes.categories, style=self.theme.tags) - if attributes.categories else None) + ( + "Categories:", + ( + smart_join(*attributes.categories, style=self.theme.tags) + if attributes.categories + else None + ), + ) ] # Updated (IP and domain) - data += [ - ("Updated:", Timestamp(attributes.last_modification_date).render) - ] + data += [("Updated:", Timestamp(attributes.last_modification_date).render)] # Last seen (Domain only) if hasattr(attributes, "last_dns_records_date"): - data += [ - ("Last Seen:", Timestamp(attributes.last_dns_records_date).render) - ] + data += [("Last Seen:", Timestamp(attributes.last_dns_records_date).render)] return self._gen_section( - self._gen_table(*data), - self._gen_heading_text("VirusTotal") + self._gen_table(*data), self._gen_heading_text("VirusTotal") ) def _gen_geoasn_section(self) -> Optional[RenderableType]: - """ IP location and ASN section. Applies to IP views only """ + """IP location and ASN section. Applies to IP views only""" enrich = self._get_geoasn_enrichment(self.entity.data.id_) data: List[Tuple[Union[str, Text], Union[RenderableType, None]]] = [] @@ -382,43 +407,46 @@ def _gen_geoasn_section(self) -> Optional[RenderableType]: ("Location:", smart_join(enrich.city, enrich.region, enrich.country)), ] return self._gen_section( - - self._gen_table(*data), - self._gen_heading_text(section_title) + self._gen_table(*data), self._gen_heading_text(section_title) ) return None # No enrichment data def _gen_shodan_section(self) -> Optional[RenderableType]: - """ Shodan section. Applies to IP views only """ + """Shodan section. Applies to IP views only""" enrich = self._get_shodan_enrichment(self.entity.data.id_) data: List[Tuple[Union[str, Text], Union[RenderableType, None]]] = [] if enrich: section_title = "Shodan" - tags = smart_join(*enrich.tags, style=self.theme.tags) if enrich.tags else None + tags = ( + smart_join(*enrich.tags, style=self.theme.tags) if enrich.tags else None + ) services_field = self._gen_linked_field_name( "Services", - hyperlink=f"{self.shodan_gui_baseurl}/{self.entity.data.id_}" + hyperlink=f"{self.shodan_gui_baseurl}/{self.entity.data.id_}", ) data += [ ("OS:", enrich.os), (services_field, self._gen_shodan_services(enrich)), ("Tags:", tags), - ("Last Scan:", Timestamp(f"{enrich.last_update}+00:00").render), # Timestamps are UTC - # (source: Google) + ( + "Last Scan:", + Timestamp(f"{enrich.last_update}+00:00").render, + ), # Timestamps are UTC + # (source: Google) ] return self._gen_section( - self._gen_table(*data), - self._gen_heading_text(section_title) + self._gen_table(*data), self._gen_heading_text(section_title) ) return None # No enrichment data def _gen_urlhaus_section(self) -> Optional[RenderableType]: - """ URLhaus """ + """URLhaus""" + def bl_text(blocklist: str, status: str) -> Text: # https://urlhaus-api.abuse.ch/#hostinfo text = Text() @@ -429,7 +457,7 @@ def bl_text(blocklist: str, status: str) -> Text: elif status.endswith("_domain") or status == "listed": text.append(status, self.theme.urlhaus_bl_high) else: # pragma: no cover - raise Exception(f"Invalid URLhaus BL status: {status}") + raise WtfisException(f"Invalid URLhaus BL status: {status}") text.append(" in ").append(blocklist, style=self.theme.urlhaus_bl_name) return text @@ -438,45 +466,66 @@ def bl_text(blocklist: str, status: str) -> Text: data: List[Tuple[Union[str, Text], Union[RenderableType, None]]] = [] if enrich: - malware_urls_field: Union[Text, str] = self._gen_linked_field_name( - "Malware URLs", - hyperlink=enrich.urlhaus_reference, - ) if enrich.urlhaus_reference else "Malware URLs:" + malware_urls_field: Union[Text, str] = ( + self._gen_linked_field_name( + "Malware URLs", + hyperlink=enrich.urlhaus_reference, + ) + if enrich.urlhaus_reference + else "Malware URLs:" + ) malware_urls_value = Text() - (malware_urls_value - .append( - (str(enrich.online_url_count) - if enrich.url_count and enrich.url_count <= 100 - else f"{enrich.online_url_count}+") + " online", - style=self.theme.error if enrich.online_url_count > 0 else self.theme.warn, - ) - .append( - f" ({enrich.url_count} total)", - style=self.theme.table_value, - )) - - tags = smart_join(*enrich.tags, style=self.theme.tags) if enrich.tags else None + ( + malware_urls_value.append( + ( + str(enrich.online_url_count) + if enrich.url_count and enrich.url_count <= 100 + else f"{enrich.online_url_count}+" + ) + + " online", + style=( + self.theme.error + if enrich.online_url_count > 0 + else self.theme.warn + ), + ).append( + f" ({enrich.url_count} total)", + style=self.theme.table_value, + ) + ) + + tags = ( + smart_join(*enrich.tags, style=self.theme.tags) if enrich.tags else None + ) data += [ (malware_urls_field, malware_urls_value), ( "Blocklists:", - (bl_text("spamhaus", enrich.blacklists.spamhaus_dbl if enrich.blacklists else "") + "\n" + - bl_text("surbl", enrich.blacklists.surbl if enrich.blacklists else "")) + ( + bl_text( + "spamhaus", + enrich.blacklists.spamhaus_dbl if enrich.blacklists else "", + ) + + "\n" + + bl_text( + "surbl", + enrich.blacklists.surbl if enrich.blacklists else "", + ) + ), ), ("Tags:", tags), ] return self._gen_section( - self._gen_table(*data), - self._gen_heading_text("URLhaus") + self._gen_table(*data), self._gen_heading_text("URLhaus") ) return None # No enrichment data def _gen_ip_other_section(self) -> Optional[RenderableType]: - """ Other section for IP views """ + """Other section for IP views""" data: List[Tuple[Union[str, Text], Union[RenderableType, None]]] = [] # Greynoise @@ -491,8 +540,7 @@ def _gen_ip_other_section(self) -> Optional[RenderableType]: if data: return self._gen_section( - self._gen_table(*data), - self._gen_heading_text("Other") + self._gen_table(*data), self._gen_heading_text("Other") ) return None # No other data @@ -507,14 +555,21 @@ def whois_panel(self) -> Optional[Panel]: else: # VT hyperlink = None - heading = self._gen_heading_text( - self.whois.domain, - hyperlink=hyperlink, - type="h2", - ) if self.whois.domain else None + heading = ( + self._gen_heading_text( + self.whois.domain, + hyperlink=hyperlink, + type="h2", + ) + if self.whois.domain + else None + ) - organization = (Text(self.whois.organization, style=self.theme.whois_org) - if self.whois.organization else None) + organization = ( + Text(self.whois.organization, style=self.theme.whois_org) + if self.whois.organization + else None + ) body = self._gen_table( ("Registrar:", self.whois.registrar), ("Organization:", organization), @@ -526,7 +581,10 @@ def whois_panel(self) -> Optional[Panel]: ("State:", self.whois.state), ("Country:", self.whois.country), ("Postcode:", self.whois.postal_code), - ("Nameservers:", smart_join(*self.whois.name_servers, style=self.theme.nameserver_list)), + ( + "Nameservers:", + smart_join(*self.whois.name_servers, style=self.theme.nameserver_list), + ), ("DNSSEC:", self.whois.dnssec), ("Registered:", Timestamp(self.whois.date_created).render), ("Updated:", Timestamp(self.whois.date_changed).render), diff --git a/wtfis/ui/view.py b/wtfis/ui/view.py index dcadcd8..446c903 100644 --- a/wtfis/ui/view.py +++ b/wtfis/ui/view.py @@ -42,16 +42,16 @@ def __init__( urlhaus: UrlHausMap, max_resolutions: int = 3, ) -> None: - super().__init__(console, entity, geoasn, whois, shodan, greynoise, abuseipdb, urlhaus) + super().__init__( + console, entity, geoasn, whois, shodan, greynoise, abuseipdb, urlhaus + ) self.resolutions = resolutions self.max_resolutions = max_resolutions def domain_panel(self) -> Panel: content = [self._gen_vt_section()] # VT section - for section in ( - self._gen_urlhaus_section(), # URLhaus section - ): + for section in (self._gen_urlhaus_section(),): # URLhaus section if section is not None: content.append("") content.append(section) @@ -71,7 +71,9 @@ def resolutions_panel(self) -> Optional[Panel]: attributes = ip.attributes # Analysis - analysis = self._gen_vt_analysis_stats(attributes.ip_address_last_analysis_stats) + analysis = self._gen_vt_analysis_stats( + attributes.ip_address_last_analysis_stats + ) analysis_field = self._gen_linked_field_name( "Analysis", hyperlink=f"{self.vt_gui_baseurl_ip}/{attributes.ip_address}", @@ -95,23 +97,30 @@ def resolutions_panel(self) -> Optional[Panel]: data += [ ("ASN:", asn), ("ISP:", geoasn.isp), - ("Location:", smart_join(geoasn.city, geoasn.region, geoasn.country)), + ( + "Location:", + smart_join(geoasn.city, geoasn.region, geoasn.country), + ), ] # Shodan shodan = self._get_shodan_enrichment(attributes.ip_address) if shodan: - tags = smart_join(*shodan.tags, style=self.theme.tags) if shodan.tags else None + tags = ( + smart_join(*shodan.tags, style=self.theme.tags) + if shodan.tags + else None + ) services_field = self._gen_linked_field_name( "Services", - hyperlink=f"{self.shodan_gui_baseurl}/{attributes.ip_address}" + hyperlink=f"{self.shodan_gui_baseurl}/{attributes.ip_address}", ) data += [ ("OS:", shodan.os), (services_field, self._gen_shodan_services(shodan)), ("Tags:", tags), - # ("Last Scan:", Timestamp(f"{shodan.last_update}+00:00").render), # Timestamps are UTC - # # (source: Google) + # Timestamps are UTC (source: Google) + # ("Last Scan:", Timestamp(f"{shodan.last_update}+00:00").render), ] # Greynoise @@ -125,12 +134,16 @@ def resolutions_panel(self) -> Optional[Panel]: data += [self._gen_abuseipdb_tuple(abuseipdb)] # Include a disclaimer if last seen is older than 1 year - # Note: Disabled for now because I originally understood that the resolution date was the last time - # the domain was resolved, but it may actually be he first time the IP itself was seen with the domain. + # Note: Disabled for now because I originally understood that the + # resolution date was the last time the domain was resolved, but it may + # actually be he first time the IP itself was seen with the domain. # if enrich and older_than(attributes.date, 365): # body = Group( # self._gen_table(*data), - # Text("**Enrichment data may be inaccurate", style=self.theme.disclaimer), + # Text( + # "**Enrichment data may be inaccurate", + # style=self.theme.disclaimer + # ), # ) # type: Union[Group, Table, str] # else: # body = self._gen_table(*data) @@ -139,10 +152,7 @@ def resolutions_panel(self) -> Optional[Panel]: content.append(self._gen_section(body, heading)) # Add extra line break if not last item in list - if ( - idx < self.max_resolutions - 1 and - idx < len(self.resolutions.data) - 1 - ): + if idx < self.max_resolutions - 1 and idx < len(self.resolutions.data) - 1: content.append("") # Info about how many more IPs were not shown @@ -151,8 +161,11 @@ def resolutions_panel(self) -> Optional[Panel]: content.append( Text(justify="center", end="\n").append( f"+{self.resolutions.meta.count - self.max_resolutions} more", - style=(f"{self.theme.footer} " - f"link {self.vt_gui_baseurl_domain}/{self.entity.data.id_}/relations"), + style=( + f"{self.theme.footer} " + f"link {self.vt_gui_baseurl_domain}/{self.entity.data.id_}" + "/relations" + ), ) ) @@ -164,11 +177,15 @@ def resolutions_panel(self) -> Optional[Panel]: return None def print(self, one_column: bool = False) -> None: - renderables = [i for i in ( - self.domain_panel(), - self.resolutions_panel(), - self.whois_panel(), - ) if i is not None] + renderables = [ + i + for i in ( + self.domain_panel(), + self.resolutions_panel(), + self.whois_panel(), + ) + if i is not None + ] if one_column: self.console.print(Group(*([""] + renderables))) # type: ignore @@ -192,15 +209,17 @@ def __init__( abuseipdb: AbuseIpDbMap, urlhaus: UrlHausMap, ) -> None: - super().__init__(console, entity, geoasn, whois, shodan, greynoise, abuseipdb, urlhaus) + super().__init__( + console, entity, geoasn, whois, shodan, greynoise, abuseipdb, urlhaus + ) def ip_panel(self) -> Panel: content = [self._gen_vt_section()] # VT section for section in ( - self._gen_geoasn_section(), # IP location and ASN section - self._gen_shodan_section(), # Shodan section - self._gen_urlhaus_section(), # URLhaus section - self._gen_ip_other_section(), # Other section + self._gen_geoasn_section(), # IP location and ASN section + self._gen_shodan_section(), # Shodan section + self._gen_urlhaus_section(), # URLhaus section + self._gen_ip_other_section(), # Other section ): if section is not None: content.append("") @@ -209,10 +228,14 @@ def ip_panel(self) -> Panel: return self._gen_panel(self._gen_group(content), self.entity.data.id_) def print(self, one_column: bool = False) -> None: - renderables = [i for i in ( - self.ip_panel(), - self.whois_panel(), - ) if i is not None] + renderables = [ + i + for i in ( + self.ip_panel(), + self.whois_panel(), + ) + if i is not None + ] if one_column: self.console.print(Group(*([""] + renderables))) # type: ignore diff --git a/wtfis/utils.py b/wtfis/utils.py index 2857066..dec6f24 100644 --- a/wtfis/utils.py +++ b/wtfis/utils.py @@ -20,7 +20,7 @@ def __init__(self, ts: Union[str, int, None]): self.timestamp = self._standardize(ts) def _standardize(self, ts: Union[str, int, None]) -> Optional[str]: - """ Convert any time to a standard format """ + """Convert any time to a standard format""" # Do nothing if ts is None if ts is None: return None @@ -63,7 +63,11 @@ def _standardize(self, ts: Union[str, int, None]) -> Optional[str]: # Default try: - return datetime.fromisoformat(ts).astimezone(timezone.utc).strftime(self.std_utc) + return ( + datetime.fromisoformat(ts) + .astimezone(timezone.utc) + .strftime(self.std_utc) + ) except ValueError: # Cannot convert; fail open return ts @@ -84,13 +88,12 @@ def render(self) -> Optional[RenderableType]: def older_than(ts: int, days: int) -> bool: - """ Tells if a timestamp is older than X days. "ts" must be epoch format """ + """Tells if a timestamp is older than X days. "ts" must be epoch format""" return datetime.fromtimestamp(ts) < datetime.now() - timedelta(days=days) def smart_join( - *items: Optional[Union[Text, str]], - style: Optional[str] = None + *items: Optional[Union[Text, str]], style: Optional[str] = None ) -> Union[Text, str]: text = Text() for idx, item in enumerate(items): @@ -110,12 +113,12 @@ def error_and_exit(message: str, status: int = 1): def refang(text: str) -> str: - """ Strip []s out of text """ + """Strip []s out of text""" return text.replace("[", "").replace("]", "") def is_ip(text: str) -> bool: - """ Detect whether text is IPv4 or not """ + """Detect whether text is IPv4 or not""" try: return ip_address(refang(text)).is_global except ValueError: From 0743add246c91f70079943bbc00913c6f43b4364 Mon Sep 17 00:00:00 2001 From: joon Date: Mon, 10 Jun 2024 02:59:45 -0700 Subject: [PATCH 3/4] Use isort --- pyproject.toml | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 54eb576..ff11332 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ dependencies = [ "flake8>=7.0.0", "flake8-bugbear~=24.4.6", "freezegun", + "isort~=5.13.2", "mypy", "pytest", "pytest-cov", @@ -90,17 +91,24 @@ dependencies = [ # Make sure the respective versions are synced with default! "black~=24.4.2", "flake8>=7.0.0", "flake8-bugbear~=24.4.6", + "isort~=5.13.2", ] [tool.hatch.envs.lint.scripts] -blackcheck = "black --check wtfis tests" +black_check = "black --check wtfis tests" +isort_check = "isort --check-only wtfis tests" flake = "flake8 wtfis" security = "bandit --quiet -r wtfis" all = [ - "blackcheck", + "black_check", + "isort_check", "flake", "security", ] +# isort +[tool.isort] +profile = "black" + # mypy [tool.mypy] ignore_missing_imports = true From 375a939757558bb574d43957cdc9d8c3bade79c9 Mon Sep 17 00:00:00 2001 From: joon Date: Mon, 10 Jun 2024 03:20:57 -0700 Subject: [PATCH 4/4] Apply isort to code --- pyproject.toml | 2 +- tests/conftest.py | 2 +- tests/test_clients.py | 2 +- tests/test_models_vt.py | 6 ++---- tests/test_ui_domain_view.py | 9 +++------ tests/test_ui_ip_view.py | 12 +++++------- tests/test_utils.py | 10 +--------- wtfis/clients/base.py | 4 ++-- wtfis/clients/greynoise.py | 3 ++- wtfis/clients/ipwhois.py | 7 +++---- wtfis/clients/shodan.py | 3 ++- wtfis/clients/types.py | 1 - wtfis/clients/virustotal.py | 7 +------ wtfis/handlers/base.py | 6 +++--- wtfis/handlers/domain.py | 4 +--- wtfis/main.py | 9 +++++---- wtfis/models/abuseipdb.py | 2 +- wtfis/models/base.py | 2 +- wtfis/models/greynoise.py | 3 ++- wtfis/models/ip2whois.py | 3 ++- wtfis/models/passivetotal.py | 3 ++- wtfis/models/shodan.py | 3 ++- wtfis/models/types.py | 1 - wtfis/models/urlhaus.py | 3 ++- wtfis/models/virustotal.py | 3 ++- wtfis/ui/base.py | 15 ++++----------- wtfis/ui/progress.py | 2 +- wtfis/ui/view.py | 19 ++++++------------- wtfis/utils.py | 1 - 29 files changed, 58 insertions(+), 89 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ff11332..1b7ea80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -96,7 +96,7 @@ dependencies = [ # Make sure the respective versions are synced with default! [tool.hatch.envs.lint.scripts] black_check = "black --check wtfis tests" isort_check = "isort --check-only wtfis tests" -flake = "flake8 wtfis" +flake = "flake8 wtfis tests" security = "bandit --quiet -r wtfis" all = [ "black_check", diff --git a/tests/conftest.py b/tests/conftest.py index 902ed8b..071e5c3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,7 @@ -import pytest from pathlib import Path from typing import Optional +import pytest from rich.console import RenderableType from rich.text import Span, Text diff --git a/tests/test_clients.py b/tests/test_clients.py index d051f5a..4eed4d2 100644 --- a/tests/test_clients.py +++ b/tests/test_clients.py @@ -1,7 +1,7 @@ import json -import pytest from unittest.mock import MagicMock, patch +import pytest from shodan import APIError from wtfis.clients.abuseipdb import AbuseIpDbClient diff --git a/tests/test_models_vt.py b/tests/test_models_vt.py index a7bffe5..217acf0 100644 --- a/tests/test_models_vt.py +++ b/tests/test_models_vt.py @@ -1,8 +1,6 @@ import json -from wtfis.models.virustotal import ( - Domain, - Resolutions, -) + +from wtfis.models.virustotal import Domain, Resolutions class TestVirustotalModels: diff --git a/tests/test_ui_domain_view.py b/tests/test_ui_domain_view.py index 03f1468..3627d9f 100644 --- a/tests/test_ui_domain_view.py +++ b/tests/test_ui_domain_view.py @@ -11,19 +11,16 @@ from wtfis.clients.greynoise import GreynoiseClient from wtfis.clients.ipwhois import IpWhoisClient from wtfis.clients.shodan import ShodanClient -from wtfis.models.abuseipdb import AbuseIpDbMap from wtfis.clients.urlhaus import UrlHausClient +from wtfis.models.abuseipdb import AbuseIpDbMap from wtfis.models.greynoise import GreynoiseIpMap from wtfis.models.ip2whois import Whois as Ip2Whois from wtfis.models.ipwhois import IpWhoisMap from wtfis.models.passivetotal import Whois as PTWhois from wtfis.models.shodan import ShodanIpMap from wtfis.models.urlhaus import UrlHausMap -from wtfis.models.virustotal import ( - Domain, - Resolutions, - Whois as VTWhois, -) +from wtfis.models.virustotal import Domain, Resolutions +from wtfis.models.virustotal import Whois as VTWhois from wtfis.ui.view import DomainView diff --git a/tests/test_ui_ip_view.py b/tests/test_ui_ip_view.py index 36cb04d..eb25f5e 100644 --- a/tests/test_ui_ip_view.py +++ b/tests/test_ui_ip_view.py @@ -1,13 +1,13 @@ -import pytest import json +from unittest.mock import MagicMock +import pytest from rich.console import Console from rich.panel import Panel from rich.table import Table from rich.text import Span, Text -from unittest.mock import MagicMock -from wtfis.clients.abuseipdb import AbuseIpDbClient +from wtfis.clients.abuseipdb import AbuseIpDbClient from wtfis.clients.greynoise import GreynoiseClient from wtfis.clients.ipwhois import IpWhoisClient from wtfis.clients.shodan import ShodanClient @@ -16,10 +16,8 @@ from wtfis.models.ipwhois import IpWhoisMap from wtfis.models.passivetotal import Whois as PTWhois from wtfis.models.urlhaus import UrlHausMap -from wtfis.models.virustotal import ( - IpAddress, - Whois as VTWhois, -) +from wtfis.models.virustotal import IpAddress +from wtfis.models.virustotal import Whois as VTWhois from wtfis.ui.view import IpAddressView diff --git a/tests/test_utils.py b/tests/test_utils.py index 8ba7b12..e4687f1 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,16 +1,8 @@ import pytest - from freezegun import freeze_time from rich.text import Text -from wtfis.utils import ( - Timestamp, - older_than, - smart_join, - error_and_exit, - refang, - is_ip, -) +from wtfis.utils import Timestamp, error_and_exit, is_ip, older_than, refang, smart_join class TestUtils: diff --git a/wtfis/clients/base.py b/wtfis/clients/base.py index 5b9c720..0d05484 100644 --- a/wtfis/clients/base.py +++ b/wtfis/clients/base.py @@ -1,9 +1,9 @@ import abc import json -import requests - from typing import Optional, Union +import requests + from wtfis.models.base import WhoisBase from wtfis.models.types import DomainEnrichmentType, IpEnrichmentType diff --git a/wtfis/clients/greynoise.py b/wtfis/clients/greynoise.py index 0cfc96a..2102f54 100644 --- a/wtfis/clients/greynoise.py +++ b/wtfis/clients/greynoise.py @@ -1,6 +1,7 @@ -from requests.exceptions import HTTPError from typing import Optional +from requests.exceptions import HTTPError + from wtfis.clients.base import BaseIpEnricherClient, BaseRequestsClient from wtfis.models.greynoise import GreynoiseIp, GreynoiseIpMap diff --git a/wtfis/clients/ipwhois.py b/wtfis/clients/ipwhois.py index 42ab8db..0fa7629 100644 --- a/wtfis/clients/ipwhois.py +++ b/wtfis/clients/ipwhois.py @@ -1,9 +1,8 @@ -from wtfis.clients.base import BaseIpEnricherClient, BaseRequestsClient -from wtfis.models.ipwhois import IpWhoisMap -from wtfis.models.ipwhois import IpWhois - from typing import Optional +from wtfis.clients.base import BaseIpEnricherClient, BaseRequestsClient +from wtfis.models.ipwhois import IpWhois, IpWhoisMap + class IpWhoisClient(BaseRequestsClient, BaseIpEnricherClient): """ diff --git a/wtfis/clients/shodan.py b/wtfis/clients/shodan.py index 3e42379..30be00b 100644 --- a/wtfis/clients/shodan.py +++ b/wtfis/clients/shodan.py @@ -1,6 +1,7 @@ -from shodan import Shodan from typing import Optional +from shodan import Shodan + from wtfis.clients.base import BaseClient, BaseIpEnricherClient from wtfis.models.shodan import ShodanIp, ShodanIpMap diff --git a/wtfis/clients/types.py b/wtfis/clients/types.py index 1226417..5115fab 100644 --- a/wtfis/clients/types.py +++ b/wtfis/clients/types.py @@ -9,7 +9,6 @@ from wtfis.clients.passivetotal import PTClient from wtfis.clients.virustotal import VTClient - # IP geolocation and ASN client types IpGeoAsnClientType = Union[IpWhoisClient,] diff --git a/wtfis/clients/virustotal.py b/wtfis/clients/virustotal.py index 4e32471..ecf45cb 100644 --- a/wtfis/clients/virustotal.py +++ b/wtfis/clients/virustotal.py @@ -1,10 +1,5 @@ from wtfis.clients.base import BaseRequestsClient, BaseWhoisClient -from wtfis.models.virustotal import ( - Domain, - IpAddress, - Resolutions, - Whois, -) +from wtfis.models.virustotal import Domain, IpAddress, Resolutions, Whois from wtfis.utils import is_ip diff --git a/wtfis/handlers/base.py b/wtfis/handlers/base.py index 5fd7674..7d01893 100644 --- a/wtfis/handlers/base.py +++ b/wtfis/handlers/base.py @@ -1,4 +1,5 @@ import abc +from typing import Callable, List, Optional, Union from pydantic import ValidationError from requests.exceptions import ( @@ -11,11 +12,11 @@ from rich.console import Console from rich.progress import Progress from shodan.exception import APIError -from typing import Callable, List, Optional, Union from wtfis.clients.abuseipdb import AbuseIpDbClient from wtfis.clients.greynoise import GreynoiseClient from wtfis.clients.shodan import ShodanClient +from wtfis.clients.types import IpGeoAsnClientType, IpWhoisClientType from wtfis.clients.urlhaus import UrlHausClient from wtfis.clients.virustotal import VTClient from wtfis.models.abuseipdb import AbuseIpDbMap @@ -23,10 +24,9 @@ from wtfis.models.greynoise import GreynoiseIpMap from wtfis.models.ipwhois import IpWhoisMap from wtfis.models.shodan import ShodanIpMap -from wtfis.models.urlhaus import UrlHausMap from wtfis.models.types import IpGeoAsnMapType +from wtfis.models.urlhaus import UrlHausMap from wtfis.models.virustotal import Domain, IpAddress -from wtfis.clients.types import IpGeoAsnClientType, IpWhoisClientType from wtfis.ui.theme import Theme from wtfis.utils import error_and_exit, refang diff --git a/wtfis/handlers/domain.py b/wtfis/handlers/domain.py index a39e126..1041fb8 100644 --- a/wtfis/handlers/domain.py +++ b/wtfis/handlers/domain.py @@ -11,17 +11,15 @@ from wtfis.clients.abuseipdb import AbuseIpDbClient from wtfis.clients.greynoise import GreynoiseClient from wtfis.clients.shodan import ShodanClient +from wtfis.clients.types import IpGeoAsnClientType, IpWhoisClientType from wtfis.clients.urlhaus import UrlHausClient from wtfis.clients.virustotal import VTClient - from wtfis.handlers.base import ( BaseHandler, common_exception_handler, failopen_exception_handler, ) - from wtfis.models.virustotal import Resolutions -from wtfis.clients.types import IpGeoAsnClientType, IpWhoisClientType class DomainHandler(BaseHandler): diff --git a/wtfis/main.py b/wtfis/main.py index fdf421c..a31d552 100644 --- a/wtfis/main.py +++ b/wtfis/main.py @@ -1,12 +1,13 @@ import argparse import os - from argparse import Namespace -from dotenv import load_dotenv from pathlib import Path +from typing import Union + +from dotenv import load_dotenv from rich.console import Console from rich.progress import Progress -from typing import Union + from wtfis.clients.abuseipdb import AbuseIpDbClient from wtfis.clients.greynoise import GreynoiseClient from wtfis.clients.ip2whois import Ip2WhoisClient @@ -20,10 +21,10 @@ from wtfis.handlers.domain import DomainHandler from wtfis.handlers.ip import IpAddressHandler from wtfis.models.virustotal import Domain, IpAddress -from wtfis.utils import error_and_exit, is_ip from wtfis.ui.base import BaseView from wtfis.ui.progress import get_progress from wtfis.ui.view import DomainView, IpAddressView +from wtfis.utils import error_and_exit, is_ip from wtfis.version import get_version diff --git a/wtfis/models/abuseipdb.py b/wtfis/models/abuseipdb.py index d79036e..a457720 100644 --- a/wtfis/models/abuseipdb.py +++ b/wtfis/models/abuseipdb.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Optional, List, Dict +from typing import Dict, List, Optional from pydantic import BaseModel, Field, RootModel diff --git a/wtfis/models/base.py b/wtfis/models/base.py index c89bce0..fb6e4d4 100644 --- a/wtfis/models/base.py +++ b/wtfis/models/base.py @@ -2,10 +2,10 @@ import abc import sys +from typing import List, Mapping, Optional from pydantic import BaseModel, BeforeValidator, ConfigDict, RootModel from pydantic.v1.validators import str_validator -from typing import List, Mapping, Optional if sys.version_info >= (3, 9): from typing import Annotated diff --git a/wtfis/models/greynoise.py b/wtfis/models/greynoise.py index 0ce83dd..c21aad7 100644 --- a/wtfis/models/greynoise.py +++ b/wtfis/models/greynoise.py @@ -1,8 +1,9 @@ from __future__ import annotations -from pydantic import BaseModel, RootModel from typing import Dict, Optional +from pydantic import BaseModel, RootModel + class GreynoiseIp(BaseModel): ip: str diff --git a/wtfis/models/ip2whois.py b/wtfis/models/ip2whois.py index 81e1f9a..84f4c9a 100644 --- a/wtfis/models/ip2whois.py +++ b/wtfis/models/ip2whois.py @@ -1,6 +1,7 @@ -from pydantic import Field, field_validator, model_validator from typing import List, Optional +from pydantic import Field, field_validator, model_validator + from wtfis.models.base import WhoisBase diff --git a/wtfis/models/passivetotal.py b/wtfis/models/passivetotal.py index 717bb2f..5dd9360 100644 --- a/wtfis/models/passivetotal.py +++ b/wtfis/models/passivetotal.py @@ -1,6 +1,7 @@ -from pydantic import Field, model_validator from typing import List, Optional +from pydantic import Field, model_validator + from wtfis.models.base import WhoisBase diff --git a/wtfis/models/shodan.py b/wtfis/models/shodan.py index 0282720..37ed960 100644 --- a/wtfis/models/shodan.py +++ b/wtfis/models/shodan.py @@ -1,9 +1,10 @@ from __future__ import annotations from collections import defaultdict, namedtuple -from pydantic import BaseModel, RootModel from typing import Dict, List, Optional +from pydantic import BaseModel, RootModel + from wtfis.models.base import LaxStr diff --git a/wtfis/models/types.py b/wtfis/models/types.py index bb3538a..b2285ef 100644 --- a/wtfis/models/types.py +++ b/wtfis/models/types.py @@ -10,7 +10,6 @@ from wtfis.models.shodan import ShodanIpMap from wtfis.models.urlhaus import UrlHausMap - # IP enrichment map types IpEnrichmentType = Union[ AbuseIpDbMap, diff --git a/wtfis/models/urlhaus.py b/wtfis/models/urlhaus.py index 20ebcda..68328ff 100644 --- a/wtfis/models/urlhaus.py +++ b/wtfis/models/urlhaus.py @@ -1,8 +1,9 @@ from __future__ import annotations -from pydantic import BaseModel, RootModel, field_validator from typing import Dict, List, Optional, Set +from pydantic import BaseModel, RootModel, field_validator + class Blacklists(BaseModel): spamhaus_dbl: str diff --git a/wtfis/models/virustotal.py b/wtfis/models/virustotal.py index e66c7f3..83d5711 100644 --- a/wtfis/models/virustotal.py +++ b/wtfis/models/virustotal.py @@ -1,6 +1,7 @@ -from pydantic import BaseModel, Field, RootModel, field_validator, model_validator from typing import Any, Dict, List, Optional +from pydantic import BaseModel, Field, RootModel, field_validator, model_validator + from wtfis.models.base import LaxStr, WhoisBase diff --git a/wtfis/ui/base.py b/wtfis/ui/base.py index 3dbb5c7..a967c3a 100644 --- a/wtfis/ui/base.py +++ b/wtfis/ui/base.py @@ -1,26 +1,19 @@ import abc +from typing import Any, Generator, List, Optional, Tuple, Union -from rich.console import ( - Console, - Group, - RenderableType, - group, -) +from rich.console import Console, Group, RenderableType, group from rich.panel import Panel from rich.table import Table from rich.text import Text -from typing import Any, Generator, List, Optional, Tuple, Union + from wtfis.exceptions import WtfisException from wtfis.models.abuseipdb import AbuseIpDb, AbuseIpDbMap from wtfis.models.base import WhoisBase from wtfis.models.greynoise import GreynoiseIp, GreynoiseIpMap from wtfis.models.shodan import ShodanIp, ShodanIpMap from wtfis.models.types import IpGeoAsnMapType, IpGeoAsnType -from wtfis.models.virustotal import ( - LastAnalysisStats, - PopularityRanks, -) from wtfis.models.urlhaus import UrlHaus, UrlHausMap +from wtfis.models.virustotal import LastAnalysisStats, PopularityRanks from wtfis.ui.theme import Theme from wtfis.utils import Timestamp, is_ip, smart_join diff --git a/wtfis/ui/progress.py b/wtfis/ui/progress.py index 329cfec..967fe52 100644 --- a/wtfis/ui/progress.py +++ b/wtfis/ui/progress.py @@ -4,8 +4,8 @@ Progress, SpinnerColumn, TaskProgressColumn, - TimeElapsedColumn, TextColumn, + TimeElapsedColumn, ) diff --git a/wtfis/ui/view.py b/wtfis/ui/view.py index 446c903..52d56a0 100644 --- a/wtfis/ui/view.py +++ b/wtfis/ui/view.py @@ -1,25 +1,18 @@ +from typing import List, Optional, Tuple, Union + from rich.columns import Columns -from rich.console import ( - Console, - Group, - RenderableType, -) +from rich.console import Console, Group, RenderableType from rich.padding import Padding from rich.panel import Panel from rich.text import Text -from typing import List, Optional, Tuple, Union -from wtfis.models.abuseipdb import AbuseIpDbMap +from wtfis.models.abuseipdb import AbuseIpDbMap from wtfis.models.base import WhoisBase from wtfis.models.greynoise import GreynoiseIpMap from wtfis.models.shodan import ShodanIpMap -from wtfis.models.urlhaus import UrlHausMap from wtfis.models.types import IpGeoAsnMapType -from wtfis.models.virustotal import ( - Domain, - IpAddress, - Resolutions, -) +from wtfis.models.urlhaus import UrlHausMap +from wtfis.models.virustotal import Domain, IpAddress, Resolutions from wtfis.ui.base import BaseView from wtfis.utils import Timestamp, smart_join diff --git a/wtfis/utils.py b/wtfis/utils.py index dec6f24..d2435ad 100644 --- a/wtfis/utils.py +++ b/wtfis/utils.py @@ -1,6 +1,5 @@ import re import sys - from datetime import datetime, timedelta, timezone from ipaddress import ip_address from typing import Optional, Union