contrib: makeseeds: Factor out ASN lookup

1 files changed, 31 insertions, 24 deletions

diff --git a/contrib/seeds/makeseeds.py b/contrib/seeds/makeseeds.py
index f5484ab33e..710881b1ae 100755
--- a/contrib/seeds/makeseeds.py
+++ b/contrib/seeds/makeseeds.py
@@ -121,6 +121,31 @@ def filtermultiport(ips):
         hist[ip['sortkey']].append(ip)
     return [value[0] for (key,value) in list(hist.items()) if len(value)==1]
 
+def lookup_asn(net, ip):
+    '''
+    Look up the asn for an IP (4 or 6) address by querying cymry.com, or None
+    if it could not be found.
+    '''
+    try:
+        if net == 'ipv4':
+            ipaddr = ip
+            prefix = '.origin'
+        else:                  # http://www.team-cymru.com/IP-ASN-mapping.html
+            res = str()                         # 2001:4860:b002:23::68
+            for nb in ip.split(':')[:4]:  # pick the first 4 nibbles
+                for c in nb.zfill(4):           # right padded with '0'
+                    res += c + '.'              # 2001 4860 b002 0023
+            ipaddr = res.rstrip('.')            # 2.0.0.1.4.8.6.0.b.0.0.2.0.0.2.3
+            prefix = '.origin6'
+
+        asn = int([x.to_text() for x in dns.resolver.query('.'.join(
+                   reversed(ipaddr.split('.'))) + prefix + '.asn.cymru.com',
+                   'TXT').response.answer][0].split('\"')[1].split(' ')[0])
+        return asn
+    except Exception:
+        sys.stderr.write('ERR: Could not resolve ASN for "' + ip + '"\n')
+        return None
+
 # Based on Greg Maxwell's seed_filter.py
 def filterbyasn(ips, max_per_asn, max_total):
     # Sift out ips by type
@@ -129,33 +154,15 @@ def filterbyasn(ips, max_per_asn, max_total):
 
     # Filter IPv46 by ASN
     result = []
-    asn_count = {}
+    asn_count = collections.defaultdict(int)
     for ip in ips_ipv46:
         if len(result) == max_total:
             break
-        try:
-            if ip['net'] == 'ipv4':
-                ipaddr = ip['ip']
-                prefix = '.origin'
-            else:                  # http://www.team-cymru.com/IP-ASN-mapping.html
-                res = str()                         # 2001:4860:b002:23::68
-                for nb in ip['ip'].split(':')[:4]:  # pick the first 4 nibbles
-                    for c in nb.zfill(4):           # right padded with '0'
-                        res += c + '.'              # 2001 4860 b002 0023
-                ipaddr = res.rstrip('.')            # 2.0.0.1.4.8.6.0.b.0.0.2.0.0.2.3
-                prefix = '.origin6'
-
-            asn = int([x.to_text() for x in dns.resolver.query('.'.join(
-                       reversed(ipaddr.split('.'))) + prefix + '.asn.cymru.com',
-                       'TXT').response.answer][0].split('\"')[1].split(' ')[0])
-            if asn not in asn_count:
-                asn_count[asn] = 0
-            if asn_count[asn] == max_per_asn:
-                continue
-            asn_count[asn] += 1
-            result.append(ip)
-        except:
-            sys.stderr.write('ERR: Could not resolve ASN for "' + ip['ip'] + '"\n')
+        asn = lookup_asn(ip['net'], ip['ip'])
+        if asn is None or asn_count[asn] == max_per_asn:
+            continue
+        asn_count[asn] += 1
+        result.append(ip)
 
     # Add back Onions
     result.extend(ips_onion)