gennet/gen_machines.py

#!/bin/python3

import ipaddress
import jinja2
import os
import sys
from pathlib import Path

class Bridge:
	def __init__(self, name, num):
		self.name = name
		self.num = num
		self.computers = []

class Computer:
	def __init__(self, name, num):
		self.name = name
		self.num = num
		self.bridges = []

	def add_to_bridge(self, bridge):
		self.bridges.append(bridge)
		bridge.computers.append(self)

	def gen_files(self):
		files = [
			'qemu.sh',
			'interfaces',
			'bird.conf',
			'hosts',
			'hostname',
			]
		context = {'c': self}
	
		loader = jinja2.FileSystemLoader('templates/')
		environment = jinja2.Environment(loader=loader, trim_blocks=True, keep_trailing_newline=True)
		for f in files:
			template = environment.get_template(f)
			(Path('output') / self.name).mkdir(parents=True,exist_ok=True)
			template.stream(context).dump('output/' + self.name + '/' + f)
		mode = os.stat('output/' + self.name + '/qemu.sh').st_mode
		os.chmod('output/' + self.name + '/qemu.sh', mode | 0o0100)

# Computers:
comps = {}
for i, n in enumerate('ABCDEFGHIX', start=1):
	comps[n] = Computer(n, i)

# Bridges
bridges = {}
for i, n in [(x, 'net_'+str(x)) for x in range(1, 8)]:
	bridges[i] = Bridge(n, i)

def dump_bridges(fn):
	with open(fn, 'w+') as f:
		for br in bridges:
			# Generate bridge.conf for the host
			f.write('allow ' + bridges[br].name + '\n')

# Connections
# dict[ net_num -> [comp_names]
conns = {
	1: 'IGH',
	2: 'HCE',
	3: 'GDF',
	4: 'ACD',
	5: 'BEF',
	6: 'ABX',
	7: 'X',
	}

print(conns)
print(comps)
print(bridges)

for net, cs in conns.items():
	for c in cs:
		comps[c].add_to_bridge(bridges[net])

os.mkdir('output')
dump_bridges('output/bridge.conf')

for c in comps.values():
	c.gen_files()