Skip to content

Commit

Permalink
Pydantic 2 compatibility
Browse files Browse the repository at this point in the history
  • Loading branch information
@pirxthepilot
pirxthepilot committed Jul 22, 2023
1 parent c11ee7c commit 16f495b
Show file tree
Hide file tree
Showing 21 changed files with 212 additions and 243 deletions.
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ classifiers = [
"Topic :: Security",
]
dependencies = [
"pydantic~=1.10.2",
"pydantic~=2.0.3",
"python-dotenv~=0.21.0",
"requests~=2.28.1",
"rich~=12.6.0",
Expand Down
6 changes: 3 additions & 3 deletions tests/conftest.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,17 +35,17 @@ def open_test_data(fname: str) -> str:

def greynoise_get(ip, pool) -> GreynoiseIp:
""" Mock replacement for GreynoiseClient().get_ip() """
return GreynoiseIp.parse_obj(pool[ip])
return GreynoiseIp.model_validate(pool[ip])


def ipwhois_get(ip, pool) -> IpWhois:
""" Mock replacement for IpWhoisClient().get_ipwhois() """
return IpWhois.parse_obj(pool[ip])
return IpWhois.model_validate(pool[ip])


def shodan_get_ip(ip, pool) -> ShodanIp:
""" Mock replacement for ShodanClient().get_ip() """
return ShodanIp.parse_obj(pool[ip])
return ShodanIp.model_validate(pool[ip])


def timestamp_text(ts) -> Optional[RenderableType]:
Expand Down
4 changes: 2 additions & 2 deletions tests/test_cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -444,7 +444,7 @@ def test_view_domain_1(self, m_domain_view, test_data):
whois_client=MagicMock(),
greynoise_client=MagicMock(),
)
entity.vt_info = Domain.parse_obj(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)
Expand All @@ -462,7 +462,7 @@ def test_view_ip_1(self, m_ip_view, test_data):
whois_client=MagicMock(),
greynoise_client=MagicMock(),
)
entity.vt_info = IpAddress.parse_obj(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)
Expand Down
20 changes: 10 additions & 10 deletions tests/test_handlers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def test_entity_refang(self, domain_handler):
def test_fetch_data_1(self, domain_handler, test_data):
""" Test with max_resolutions = 3 (default) """
handler = domain_handler()
handler.resolutions = Resolutions.parse_obj(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()
Expand Down Expand Up @@ -135,9 +135,9 @@ def test_vt_validation_error(self, mock_requests_get, domain_handler, capsys):

capture = capsys.readouterr()

assert capture.err == (
assert capture.err.startswith(
"Data model validation error: 1 validation error for Domain\ndata\n"
" field required (type=value_error.missing)\n"
" Field required [type=missing, input_value={'intentionally': 'wrong data'}, input_type=dict]\n"
)
assert e.type == SystemExit
assert e.value.code == 1
Expand All @@ -150,7 +150,7 @@ def test_vt_validation_error(self, mock_requests_get, domain_handler, capsys):
@patch.object(requests.Session, "get")
def test_ipwhois_http_error(self, mock_requests_get, domain_handler, capsys, test_data):
handler = domain_handler()
handler.resolutions = Resolutions.parse_obj(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()

Expand Down Expand Up @@ -231,7 +231,7 @@ 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.parse_obj(json.loads(test_data("vt_resolutions_gist.json")))
handler.resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json")))

handler._fetch_greynoise()
assert handler.warnings[0].startswith("Could not fetch Greynoise: 429 Client Error:")
Expand All @@ -248,14 +248,14 @@ def test_greynoise_404_error(self, mock_requests_get, domain_handler, test_data)
handler = domain_handler(3)
mock_resp = requests.models.Response()

handler.resolutions = Resolutions.parse_obj(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

handler._fetch_greynoise()
assert len(handler.warnings) == 0
assert handler.greynoise == GreynoiseIpMap(__root__={})
assert handler.greynoise == GreynoiseIpMap.model_validate({})


class TestIpAddressHandler:
Expand Down Expand Up @@ -310,9 +310,9 @@ def test_vt_validation_error(self, mock_requests_get, ip_handler, capsys):

capture = capsys.readouterr()

assert capture.err == (
assert capture.err.startswith(
"Data model validation error: 1 validation error for IpAddress\ndata\n"
" field required (type=value_error.missing)\n"
" Field required [type=missing, input_value={'intentionally': 'wrong data'}, input_type=dict]\n"
)
assert e.type == SystemExit
assert e.value.code == 1
Expand Down Expand Up @@ -403,4 +403,4 @@ def test_greynoise_404_error(self, mock_requests_get, ip_handler):

handler._fetch_greynoise()
assert len(handler.warnings) == 0
assert handler.greynoise == GreynoiseIpMap(__root__={})
assert handler.greynoise == GreynoiseIpMap.model_validate({})
6 changes: 3 additions & 3 deletions tests/test_models_vt.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,21 @@

class TestVirustotalModels:
def test_domain_1(self, test_data):
domain = Domain.parse_obj(json.loads(test_data("vt_domain_gist.json")))
domain = Domain.model_validate(json.loads(test_data("vt_domain_gist.json")))

assert domain.data.id_ == "gist.github.com"
assert domain.data.type_ == "domain"
assert domain.data.attributes.last_analysis_stats.malicious == 0
assert domain.data.attributes.last_analysis_stats.suspicious == 0
assert domain.data.attributes.last_analysis_stats.harmless == 84
assert domain.data.attributes.last_analysis_stats.undetected == 10
assert list(domain.data.attributes.last_analysis_results.__root__.keys())[:2] == [
assert list(domain.data.attributes.last_analysis_results.root.keys())[:2] == [
"CMC Threat Intelligence",
"Snort IP sample list",
]

def test_resolutions_1(self, test_data):
res = Resolutions.parse_obj(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
Expand Down
46 changes: 23 additions & 23 deletions tests/test_ui_domain_view.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
@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.parse_obj(json.loads(test_data("vt_resolutions_gist.json")))
resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json")))

ipwhois_pool = json.loads(test_data("ipwhois_gist.json"))
ipwhois_client = IpWhoisClient()
Expand All @@ -34,11 +34,11 @@ def view01(test_data, mock_ipwhois_get):

return DomainView(
console=Console(),
entity=Domain.parse_obj(json.loads(test_data("vt_domain_gist.json"))),
entity=Domain.model_validate(json.loads(test_data("vt_domain_gist.json"))),
resolutions=resolutions,
whois=PTWhois.parse_obj(json.loads(test_data("pt_whois_gist.json"))),
whois=PTWhois.model_validate(json.loads(test_data("pt_whois_gist.json"))),
ip_enrich=ip_enrich,
greynoise=GreynoiseIpMap(__root__={}),
greynoise=GreynoiseIpMap.model_validate({}),
)


Expand All @@ -51,10 +51,10 @@ def view02(test_data):
return DomainView(
console=Console(),
entity=MagicMock(),
resolutions=Resolutions.parse_obj(json.loads(test_data("vt_resolutions_gist.json"))),
whois=VTWhois.parse_obj(json.loads(test_data("vt_whois_gist.json"))),
ip_enrich=IpWhoisMap(__root__={}),
greynoise=GreynoiseIpMap(__root__={}),
resolutions=Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json"))),
whois=VTWhois.model_validate(json.loads(test_data("vt_whois_gist.json"))),
ip_enrich=IpWhoisMap.model_validate({}),
greynoise=GreynoiseIpMap.model_validate({}),
max_resolutions=1,
)

Expand All @@ -66,7 +66,7 @@ def view03(test_data):
console=Console(),
entity=MagicMock(),
resolutions=MagicMock(),
whois=VTWhois.parse_obj(json.loads(test_data("vt_whois_bbc.json"))),
whois=VTWhois.model_validate(json.loads(test_data("vt_whois_bbc.json"))),
ip_enrich=MagicMock(),
greynoise=MagicMock(),
)
Expand All @@ -80,7 +80,7 @@ def view04(test_data):
"""
return DomainView(
console=Console(),
entity=Domain.parse_obj(json.loads(test_data("vt_domain_google.json"))),
entity=Domain.model_validate(json.loads(test_data("vt_domain_google.json"))),
resolutions=None,
whois=MagicMock(),
ip_enrich=MagicMock(),
Expand All @@ -93,7 +93,7 @@ def view05(test_data):
""" tucows.com domain. Domain test only. Test domain with negative reputation and no popularity."""
return DomainView(
console=Console(),
entity=Domain.parse_obj(json.loads(test_data("vt_domain_tucows.json"))),
entity=Domain.model_validate(json.loads(test_data("vt_domain_tucows.json"))),
resolutions=MagicMock(),
whois=MagicMock(),
ip_enrich=MagicMock(),
Expand All @@ -108,7 +108,7 @@ def view06(test_data):
console=Console(),
entity=MagicMock(),
resolutions=MagicMock(),
whois=VTWhois.parse_obj(json.loads(test_data("vt_whois_example_2.json"))),
whois=VTWhois.model_validate(json.loads(test_data("vt_whois_example_2.json"))),
ip_enrich=MagicMock(),
greynoise=MagicMock(),
)
Expand All @@ -117,7 +117,7 @@ def view06(test_data):
@pytest.fixture()
def view07(test_data, mock_shodan_get_ip):
""" gist.github.com with Shodan. Only test resolution and IP enrich. """
resolutions = Resolutions.parse_obj(json.loads(test_data("vt_resolutions_gist.json")))
resolutions = Resolutions.model_validate(json.loads(test_data("vt_resolutions_gist.json")))

shodan_pool = json.loads(test_data("shodan_gist.json"))
shodan_client = ShodanClient(MagicMock())
Expand All @@ -130,14 +130,14 @@ def view07(test_data, mock_shodan_get_ip):
resolutions=resolutions,
whois=MagicMock(),
ip_enrich=ip_enrich,
greynoise=GreynoiseIpMap(__root__={}),
greynoise=GreynoiseIpMap.model_validate({}),
)


@pytest.fixture()
def view08(test_data, mock_shodan_get_ip):
""" www.wired.com with Shodan. Only test resolution and IP enrich. """
resolutions = Resolutions.parse_obj(json.loads(test_data("vt_resolutions_wired.json")))
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())
Expand All @@ -150,15 +150,15 @@ def view08(test_data, mock_shodan_get_ip):
resolutions=resolutions,
whois=MagicMock(),
ip_enrich=ip_enrich,
greynoise=GreynoiseIpMap(__root__={}),
greynoise=GreynoiseIpMap.model_validate({}),
max_resolutions=1,
)


@pytest.fixture()
def view09(test_data, mock_shodan_get_ip, mock_greynoise_get):
""" one.one.one.one with Shodan. Only test resolution and IP enrich. """
resolutions = Resolutions.parse_obj(json.loads(test_data("vt_resolutions_one.json")))
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())
Expand Down Expand Up @@ -188,16 +188,16 @@ def view10(test_data):
console=Console(),
entity=MagicMock(),
resolutions=MagicMock(),
whois=VTWhois.parse_obj(json.loads(test_data("vt_whois_foo.json"))),
whois=VTWhois.model_validate(json.loads(test_data("vt_whois_foo.json"))),
ip_enrich=MagicMock(),
greynoise=GreynoiseIpMap(__root__={}),
greynoise=GreynoiseIpMap.model_validate({}),
)


@pytest.fixture()
def view11(test_data, mock_shodan_get_ip):
""" gist.github.com with Shodan. Only test IP enrich. Test empty open ports. """
resolutions = Resolutions.parse_obj(json.loads(test_data("vt_resolutions_gist.json")))
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())
Expand All @@ -210,7 +210,7 @@ def view11(test_data, mock_shodan_get_ip):
resolutions=resolutions,
whois=MagicMock(),
ip_enrich=ip_enrich,
greynoise=GreynoiseIpMap(__root__={}),
greynoise=GreynoiseIpMap.model_validate({}),
)


Expand All @@ -221,7 +221,7 @@ def view12(test_data):
console=Console(),
entity=MagicMock(),
resolutions=MagicMock(),
whois=Ip2Whois.parse_obj(json.loads(test_data("ip2whois_whois_hotmail.json"))),
whois=Ip2Whois.model_validate(json.loads(test_data("ip2whois_whois_hotmail.json"))),
ip_enrich=MagicMock(),
greynoise=MagicMock(),
)
Expand All @@ -234,7 +234,7 @@ def view13(test_data):
console=Console(),
entity=MagicMock(),
resolutions=MagicMock(),
whois=Ip2Whois.parse_obj(json.loads(test_data("ip2whois_whois_bbc.json"))),
whois=Ip2Whois.model_validate(json.loads(test_data("ip2whois_whois_bbc.json"))),
ip_enrich=MagicMock(),
greynoise=MagicMock(),
)
Expand Down
22 changes: 11 additions & 11 deletions tests/test_ui_ip_view.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,8 @@ def view01(test_data, mock_ipwhois_get, mock_greynoise_get):

return IpAddressView(
console=Console(),
entity=IpAddress.parse_obj(json.loads(test_data("vt_ip_1.1.1.1.json"))),
whois=PTWhois.parse_obj(json.loads(test_data("pt_whois_1.1.1.1.json"))),
entity=IpAddress.model_validate(json.loads(test_data("vt_ip_1.1.1.1.json"))),
whois=PTWhois.model_validate(json.loads(test_data("pt_whois_1.1.1.1.json"))),
ip_enrich=ip_enrich,
greynoise=greynoise_enrich,
)
Expand All @@ -59,7 +59,7 @@ def view02(test_data, mock_shodan_get_ip, mock_greynoise_get):

return IpAddressView(
console=Console(),
entity=IpAddress.parse_obj(json.loads(test_data("vt_ip_1.1.1.1.json"))),
entity=IpAddress.model_validate(json.loads(test_data("vt_ip_1.1.1.1.json"))),
whois=MagicMock(),
ip_enrich=ip_enrich,
greynoise=greynoise_enrich,
Expand All @@ -72,7 +72,7 @@ def view03(test_data):
return IpAddressView(
console=Console(),
entity=MagicMock(),
whois=VTWhois.parse_obj(json.loads(test_data("vt_whois_1.1.1.1.json"))),
whois=VTWhois.model_validate(json.loads(test_data("vt_whois_1.1.1.1.json"))),
ip_enrich=MagicMock(),
greynoise=MagicMock(),
)
Expand All @@ -86,10 +86,10 @@ def view04(test_data):
"""
return IpAddressView(
console=Console(),
entity=IpAddress.parse_obj(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"))),
whois=MagicMock(),
ip_enrich=IpWhoisMap(__root__={}),
greynoise=GreynoiseIpMap(__root__={}),
ip_enrich=IpWhoisMap.model_validate({}),
greynoise=GreynoiseIpMap.model_validate({}),
)


Expand All @@ -104,9 +104,9 @@ def view05(test_data, mock_greynoise_get):

return IpAddressView(
console=Console(),
entity=IpAddress.parse_obj(json.loads(test_data("vt_ip_1.1.1.1.json"))),
entity=IpAddress.model_validate(json.loads(test_data("vt_ip_1.1.1.1.json"))),
whois=MagicMock(),
ip_enrich=IpWhoisMap(__root__={}),
ip_enrich=IpWhoisMap.model_validate({}),
greynoise=greynoise_enrich,
)

Expand All @@ -122,9 +122,9 @@ def view06(test_data, mock_greynoise_get):

return IpAddressView(
console=Console(),
entity=IpAddress.parse_obj(json.loads(test_data("vt_ip_1.1.1.1.json"))),
entity=IpAddress.model_validate(json.loads(test_data("vt_ip_1.1.1.1.json"))),
whois=MagicMock(),
ip_enrich=IpWhoisMap(__root__={}),
ip_enrich=IpWhoisMap.model_validate({}),
greynoise=greynoise_enrich,
)

Expand Down
Loading

0 comments on commit 16f495b

Please sign in to comment.