import sys
import os.path
import os
import urllib
#http://www.icir.org/robin/pysubnettree/ 
# > wget http://www.icir.org/robin/pysubnettree/pysubnettree-0.12.tar.gz
# > tar xzf pysubnettree-0.12.tar.gz
# > cd pysubnettree-0.12
# > python setup.py install
import SubnetTree  # used for longest prefix match

def createPrefixTree( prefix_to_as ):
	print "Creating prefix to AS mapping..."
	prefix = SubnetTree.SubnetTree()

	file = open( prefix_to_as, "r" )
	while 1:
		lines = file.readlines(10000)
		if not lines:
			break
		for line in lines:
			row = line.split(' ')
			try:
				asn = min(row[1].rstrip().split("_"))  # take the minimum in case of AS sets
				pref = row[0].rstrip()
				prefix[pref]=(asn,pref)
			except:
				print "Skipped line ",line," due to format error"
				
	print "done."
	return prefix


def download(url, force=False, localName = None, localPath = "./data/"):
	"""Copy the contents of a file from a given URL
	to a local file.
	"""
	
	if ( localName== None ):
		localName = url.split('/')[-1]
	if ( force==True or not os.path.exists(localPath+localName) ):
		print "Downloading "+url+"..."
		webFile = urllib.urlopen(url)
		localFile = open(localPath+localName, 'w')
		localFile.write(webFile.read())
		webFile.close()
		localFile.close()
		print "done."
		
	return localPath+localName

def main(argv):
	if len(argv)==0 or argv[0]=="-h" or  argv[0]=="--h":
		usage()
	
	else:
		data = "./"
		# check if use want to change dest
		if argv[0]=="-d" or  argv[0]=="-dir":
			data = argv[1]
			argv = argv[2:]
			
		# need to download the files into local, and use them
		#prefix_to_as = download("http://iplane.cs.washington.edu/data/origin_as_mapping.txt")
		prefix_to_as = download("http://iplane.cs.washington.edu/data/bgp_feeds/asmapping.txt")
		if ( prefix_to_as == None ):
			prefix_to_as = download("http://iplane.cs.washington.edu/data/origin_as_mapping.txt")
			
		# ok, now we have the files and we can use them. let's open
		if ( prefix_to_as != None ):
			# create prefix tree
			prefixes = createPrefixTree( prefix_to_as )
			
			for ip in argv:
				a=prefixes[ip]
				print "%s,%s,%s"%(ip,a[0],a[1])
					
		else:
			print "failed downloading prefix to AS mapping, or ./data directory does not exist?"
		

def usage():
	print "Resolves an IP to AS number."
	print "getasn.py [-d data_dir] ip1 ip2 ip3..."

# MAIN
if __name__ == "__main__":
	main(sys.argv[1:])