async def test_alias_a_record_resolves(): # The alias name is an invalid host and so should never resolve results = await RecordResolver({}, False).resolve({ "Name": "a.dsd.io.", "Type": "A", "AliasTarget": { "HostedZoneId": "Z1BKCTXD74EZPE", "DNSName": "localhost", "EvaluateTargetHealth": False } }) print(results) assert len(results) <= 2 assert HostToScan("127.0.0.1", "a.dsd.io.") in results if len(results) == 2: assert HostToScan("::1", "a.dsd.io.") in results
async def test_ip4_a_record(): results = await RecordResolver({}, False).resolve({ "Name": "a.dsd.io.", "Type": "A", "ResourceRecords": [{ "Value": "123.321.123.1" }, { "Value": "123.321.123.2" }] }) assert results == [ HostToScan("123.321.123.1", "a.dsd.io."), HostToScan("123.321.123.2", "a.dsd.io.") ]
async def test_ip6_a_record(): results = await RecordResolver({}, False).resolve({ "Name": "b.dsd.io.", "Type": "AAAA", "ResourceRecords": [{ "Value": "2001:0db8:85a3:0000:0000:8a2e:0370:7334" }, { "Value": "9999:0db8:85a3:0000:0000:8a2e:0370:7334" }] }) assert results == [ HostToScan("2001:0db8:85a3:0000:0000:8a2e:0370:7334", "b.dsd.io."), HostToScan("9999:0db8:85a3:0000:0000:8a2e:0370:7334", "b.dsd.io.") ]
async def _resolve_a_record(self, record): if "AliasTarget" in record: return await self._resolve_alias_record(record) else: # print(f"Ingested {record['Type']} record: {record['Name']}") return self._ips_resolved([ HostToScan(resource["Value"], record["Name"]) for resource in record["ResourceRecords"] ])
async def test_alias_cname_record_resolves(): # The alias name is an invalid host and so should never resolve results = await RecordResolver({}, False).resolve({ "Name": "a.dsd.io.", "Type": "CNAME", "ResourceRecords": [ { "Value": "localhost" }, # This is yuck, I should be using patching to mock out the nslookup # instead I am relying on the build machine having internet access # and being able to resolve google { "Value": "www.google.com" } ] }) assert len(results) > 1 assert HostToScan("127.0.0.1", "a.dsd.io.") in results if HostToScan("::1", "a.dsd.io.") in results: assert len(results) > 2
async def _resolve_using_ns_lookup(self, record, redirects, record_type): # print(f"Ingested {record_type} record: {record['Name']}") to_scan = [] for redirect_to in redirects: try: to_scan += [ HostToScan(ip, record["Name"]) for ip in await ns_lookup(redirect_to) ] # print(f"Resolved ip addresses of {redirect_to} a(n) {record_type} of {record['Name']}") except gaierror: # print(f"Unable to resolve {redirect_to} a(n) {record_type} of {record['Name']}") pass return self._ips_resolved(to_scan)
async def test_uses_schedule(): schedule = MagicMock() scan_plan_table = mock_table() host_table = mock_table() address_table = mock_table() address_info_table = mock_table() writer = PlannedScanDbWriter(scan_plan_table, host_table, address_table, address_info_table, 999, schedule) await writer.prepare() # asyncly call write 10 times await gather(*[ writer.write(HostToScan(f"123.4.5.{x}", "foo.bar")) for x in range(0, 10) ]) await writer.commit() assert schedule.__next__.call_count == 10
async def test_dynamo_db_call_params(): schedule = MagicMock() schedule.__next__.return_value = 1050 scan_plan_table = mock_table() host_table = mock_table() address_table = mock_table() address_info_table = mock_table() writer = PlannedScanDbWriter(scan_plan_table, host_table, address_table, address_info_table, 999, schedule) await writer.prepare() await writer.write(HostToScan("123.4.5.6", "foo.bar")) await writer.commit() scan_plan_table.update_item.assert_called_once_with( Key={"Address": "123.4.5.6"}, UpdateExpression= "SET PlannedScanTime = :PlannedScanTime, DnsIngestTime = :DnsIngestTime", ExpressionAttributeValues={ ":PlannedScanTime": 1050, ":DnsIngestTime": 999 })
async def test_is_async(): schedule = MagicMock() scan_plan_table = mock_table() host_table = mock_table() address_table = mock_table() address_info_table = mock_table() writer = PlannedScanDbWriter(scan_plan_table, host_table, address_table, address_info_table, 999, schedule) await writer.prepare() # asyncly call write 10 times all_writes = gather(*[ writer.write(HostToScan(f"123.4.5.{x}", "foo.bar")) for x in range(0, 10) ]) assert not all_writes.done() # asyncly set the results await all_writes await writer.commit() # check that we now have the calls made assert scan_plan_table.update_item.call_count == 10