Write to exclusions.conf

Adds a main() and calls with idiomatic `if __name__` block.
Some PEP-8 refactoring.

avoidr/avoidr.py

@@ -12,6 +12,7 @@ from zipfile import ZipFile
 grand_total = {'4': 0, '6': 0}
 results     = dict()
 
+
 def calculate_hash(path):
 	''' Calculate the SHA1 hash of a file. '''
 	hash_sha1 = hashlib.sha1()
@@ -20,6 +21,7 @@ def calculate_hash(path):
 			hash_sha1.update(chunk)
 	return hash_sha1.hexdigest()
 
+
 def get_url(url, git=False) -> str:
 	''' Get the contents of a URL. '''
 	data = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_3) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/35.0.1916.47 Safari/537.36'}
@@ -28,6 +30,7 @@ def get_url(url, git=False) -> str:
 	req = urllib.request.Request(url, headers=data)
 	return urllib.request.urlopen(req, timeout=10).read().decode()
 
+
 def update_database():
 	''' Update the ASN database.  '''
 	DB = 'databases/fullASN.json.zip'
@@ -54,6 +57,7 @@ def update_database():
 		with ZipFile(DB) as zObject:
 			zObject.extract(DB[10:-4], 'databases')
 
+
 def process_asn(data):
 	''' Process an ASN. '''
 	if data['asn'] not in results:
@@ -70,27 +74,55 @@ def process_asn(data):
 			grand_total['6'] += total
 			print('Found \033[93mAS{0}\033[0m \033[1;30m({1})\033[0m containing \033[32m{2:,}\033[0m IPv6 ranges with \033[36m{3:,}\033[0m total IP addresses'.format(data['asn'], title, len(data['prefixesIPv6']), total))
 
+
 def total_ips(ranges, total=0):
 	for _range in ranges:
 		total += ipaddress.ip_network(_range).num_addresses
 	return total
 
-# Main
-print('[~] Checking for database updates...')
-update_database()
-data    = json.loads(open('databases/fullASN.json').read())
-queries = [item.rstrip() for item in open('custom.txt').readlines()]
-print('[~] Searching {len(queries):,} queries against {len(data):,} ASNs...')
-for item in data:
-	for field in [x for x in data[item] if x in ('descr','org')]:
-		if [x for x in queries if x.lower() in data[item][field].lower()]:
-			process_asn(data[item])
-			break
-with open('out.json', 'w') as fp:
-	json.dump(results, fp)
-total_v4 = ipaddress.ip_network('0.0.0.0/0').num_addresses
-total_v6 = ipaddress.ip_network('::/0').num_addresses
-print('Total IPv4 Addresses   : {0:,}'.format(total_v4))
-print('Total IPv4 After Clean : {0:,}'.format(total_v4-grand_total['4']))
-print('Total IPv6 Addresses   : {0:,}'.format(total_v6))
-print('Total IPv6 After Clean : {0:,}'.format(total_v6-grand_total['6']))
+
+def write_exclusions(asn_list, exclusions_file='exclusions.conf'):
+	''' Write masscan exclusions file. '''
+	exclusions = []
+	for asn in asn_list:
+		name = asn_list[asn]['name']
+		ranges = asn_list[asn]['ranges']
+
+		exclusions.append(f'# {name}')
+		for proto in ranges:
+			exclusions.extend(ranges[proto])
+
+	with open(exclusions_file, 'w') as fp:
+		fp.write('\n'.join(exclusions))
+
+
+def main():
+	print('[~] Checking for database updates...')
+	update_database()
+
+	data    = json.loads(open('databases/fullASN.json').read())
+	queries = [item.rstrip() for item in open('custom.txt').readlines()]
+
+	print('[~] Searching {len(queries):,} queries against {len(data):,} ASNs...')
+	for item in data:
+		for field in [x for x in data[item] if x in ('descr','org')]:
+			if [x for x in queries if x.lower() in data[item][field].lower()]:
+				process_asn(data[item])
+				break
+
+	with open('out.json', 'w') as fp:
+		json.dump(results, fp)
+
+	write_exclusions(results)
+
+	total_v4 = ipaddress.ip_network('0.0.0.0/0').num_addresses
+	total_v6 = ipaddress.ip_network('::/0').num_addresses
+	print('Total IPv4 Addresses   : {0:,}'.format(total_v4))
+	print('Total IPv4 After Clean : {0:,}'.format(total_v4-grand_total['4']))
+	print('Total IPv6 Addresses   : {0:,}'.format(total_v6))
+	print('Total IPv6 After Clean : {0:,}'.format(total_v6-grand_total['6']))
+	print('Exclusions written to exclusions.conf')
+
+
+if __name__ == '__main__':
+	main()