# This script can be used to generate config files and setup the
# directory structure for a private network of Tor nodes
# 
# Try python make-private-tor-network.py --help
# Adam Langley <agl@imperialviolet.org>

import subprocess
import os
import sys

def run(binary, *args):
	"""Run a process with the given arguments (which should include the
	 binary name) and return the output sent to stdout"""
	process = subprocess.Popen(args, executable = binary, stdout = subprocess.PIPE, stderr = None, stdin = None)

	output = []
	while True:
		data = process.stdout.read()
		if not len(data):
			break
		output.append(data)
	process.wait()
	return ''.join(output)

def tor_get_version(binary):
	"Get the version of a tor binary by running with --version"

	output = run(binary, "tor", "--version")
	lines = output.split("\n")
	# last line should look like "Tor version xxx." but it's followed by a blank line
	last = lines[-2]
	if not last.startswith("Tor version "):
		print "Failed to get version from Tor binary. Output was:"
		print
		print output
		sys.exit(1)
	return last[12:-1]

def tor_get_fingerprint(binary, torrc_path):
	output = run(binary, "tor", "-f", torrc_path, "--list-fingerprint")
	lines = output.split("\n")
	last = lines[-2]
	# the fingerprint line looks like:
	# <nickname> fingerprint
	fpr = last[last.index(" ") + 1:]
	for char in fpr:
		if (char < '0' or char > '9') and (char < 'A' or char > 'F') and char != ' ':
			print "Failed to get fingerprint from node. Output was:"
			print
			print output
			sys.exit(1)
	return fpr

def torrc_write(file, socksport, orport, dirport, nickname, basepath, dirservers):
	file.write("SocksPort %d\n" % socksport)
	file.write("SocksBindAddress 127.0.0.1\n")
	file.write("AllowUnverifiedNodes middle,rendezvous\n")
	file.write("Log debug file %s\n" % (os.path.join(basepath, "debug.log")))
	file.write("Log notice file %s\n" % (os.path.join(basepath, "log")))
	file.write("DataDirectory %s\n" % (os.path.join(basepath, "data")))
	file.write("Nickname %s\n" % nickname)
	file.write("Address 127.0.0.1\n")
	file.write("ContactInfo Private node\n")
	file.write("ORPort %d\n" % orport)
	file.write("ORBindAddress 127.0.0.1\n")
	file.write("DirPort %d\n" % dirport)
	file.write("DirBindAddress 127.0.0.1\n")
	file.write("AssumeReachable 1\n")
	file.write("DirAllowPrivateAddresses 1\n")
	
	for x in dirservers:
		file.write("DirServer %s\n" % x)

def main(options):
	tor_version = tor_get_version(options.binary)

	if options.extraconfig:
		extra_config = file(options.extraconfig, 'r').read()
	else:
		extra_config = ""

	dirservers = []
	approved_routers = []

	def mkdir(path):
		try:
			os.mkdir(path)
		except OSError:
			pass

	for config_pass in [0, 1]:
		current_port = options.baseport
		for x in range(options.num_dirservers + options.num_nodes):
			directory_server = x < options.num_dirservers
			if directory_server:
				nickname = options.basename + 'D' + str(x)
			else:
				nickname = options.basename + 'N' + str(x)
			nodebase = os.path.join(options.basepath, nickname)

			if config_pass == 0:
				mkdir(nodebase)
				mkdir(os.path.join(nodebase, 'data'))
				if directory_server:
					file(os.path.join(nodebase, 'data', 'approved-routers'), 'w+').close()
			else:
				if directory_server:
					f = file(os.path.join(nodebase, 'data', 'approved-routers'), 'w+')
					for router in approved_routers:
						f.write(router)
						f.write("\n")
					f.close()

			rc = file(os.path.join(nodebase, 'torrc'), 'w+')
			if config_pass == 0:
				listed_dirservers = ['127.0.0.1:1 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000']
			else:
				listed_dirservers = dirservers
			torrc_write(rc, orport = current_port, dirport = current_port + 1, socksport = current_port + 2, nickname = nickname, basepath = nodebase, dirservers = listed_dirservers)
			if directory_server:
				rc.write("AuthoritativeDirectory 1\n")
				rc.write("RecommendedVersions %s\n" % tor_version)
			if config_pass == 1:
				rc.write(extra_config)
			rc.close()

			if config_pass == 0:
				fingerprint = tor_get_fingerprint(options.binary, os.path.join(nodebase, 'torrc'))
				if directory_server:
					dirservers.append("127.0.0.1:%d %s" % (current_port + 1, fingerprint))
				approved_routers.append("%s %s" % (nickname, fingerprint))

			if config_pass == 1:
				print "%s -f %s &  # or: %d, dir: %d, socks: %d" % (options.binary, os.path.join(nodebase, 'torrc'), current_port, current_port + 1, current_port + 2)

			current_port += 3

if __name__ == '__main__':
	from optparse import OptionParser
	parser = OptionParser()
	parser.add_option('-b', '--tor-binary', dest = 'binary', default = 'tor', help = 'location of the tor binary')
	parser.add_option('-n', '--base-name', dest = 'basename', default = 'node', help = 'base name of the directories and nicknames of the nodes')
	parser.add_option('', '--base-port', type = 'int', dest = 'baseport', default = '3000', help = 'lower most port number to start assigning from')
	parser.add_option('-d', '--num-dirservers', type = 'int', dest = 'num_dirservers', default = '1', help = 'number of dirservers')
	parser.add_option('-o', '--num-nodes', type = 'int', dest = 'num_nodes', default = '3', help = 'number of non-dirserver nodes')
	parser.add_option('-e', '--extra-config', dest = 'extraconfig', default = None, help = 'location of a file which is appended onto the generated configs')
	parser.add_option('-p', '--path', dest = 'basepath', default = '.', help = 'the path in which to create the directories')

	(options, args) = parser.parse_args()

	if len(args):
		print "Unknown trailing arguments given"
	else:
		main(options)