Add support for getting country codes for ASNs.

This commit is contained in:
Jon Chambers 2021-05-17 19:03:04 -04:00 committed by Jon Chambers
parent f8c623074b
commit b89de860d3
2 changed files with 22 additions and 7 deletions

View File

@ -10,6 +10,8 @@ import java.io.IOException;
import java.io.Reader;
import java.net.Inet4Address;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.NavigableMap;
import java.util.Optional;
import java.util.TreeMap;
@ -21,7 +23,8 @@ import org.apache.commons.csv.CSVRecord;
* Allows IP->ASN lookup operations using data from https://iptoasn.com/.
*/
class AsnTable {
private final NavigableMap<Long, AsnRange> asnBlocksByFirstIp;
private final NavigableMap<Long, AsnRange> asnBlocksByFirstIp = new TreeMap<>();
private final Map<Long, String> countryCodesByAsn = new HashMap<>();
private static class AsnRange {
private final long rangeStart;
@ -47,23 +50,20 @@ class AsnTable {
public static final AsnTable EMPTY = new AsnTable();
public AsnTable(final Reader tsvReader) throws IOException {
final TreeMap<Long, AsnRange> treeMap = new TreeMap<>();
try (final CSVParser csvParser = CSVFormat.TDF.parse(tsvReader)) {
for (final CSVRecord record : csvParser) {
final long start = Long.parseLong(record.get(0), 10);
final long end = Long.parseLong(record.get(1), 10);
final long asn = Long.parseLong(record.get(2), 10);
final String countryCode = record.get(3);
treeMap.put(start, new AsnRange(start, end, asn));
asnBlocksByFirstIp.put(start, new AsnRange(start, end, asn));
countryCodesByAsn.put(asn, countryCode);
}
}
asnBlocksByFirstIp = treeMap;
}
private AsnTable() {
asnBlocksByFirstIp = new TreeMap<>();
}
public Optional<Long> getAsn(final Inet4Address address) {
@ -74,6 +74,10 @@ class AsnTable {
.map(entry -> entry.getValue().getAsn());
}
public Optional<String> getCountryCode(final long asn) {
return Optional.ofNullable(countryCodesByAsn.get(asn));
}
@VisibleForTesting
static long ipToLong(final Inet4Address address) {
final ByteBuffer buffer = ByteBuffer.allocate(8);

View File

@ -30,6 +30,17 @@ class AsnTableTest {
}
}
@Test
void getCountryCode() throws IOException {
try (final InputStreamReader reader = new InputStreamReader(getClass().getResourceAsStream("ip2asn-test.tsv"))) {
final AsnTable asnTable = new AsnTable(reader);
assertEquals(Optional.of("US"), asnTable.getCountryCode(7922));
assertEquals(Optional.of("VN"), asnTable.getCountryCode(7552));
assertEquals(Optional.empty(), asnTable.getCountryCode(1234));
}
}
@Test
void ipToLong() throws UnknownHostException {
assertEquals(0x00000000ffffffffL, AsnTable.ipToLong((Inet4Address) Inet4Address.getByName("255.255.255.255")));