diff --git a/src/wg_builder/builder.py b/src/wg_builder/builder.py index bbd18da..68c3eee 100644 --- a/src/wg_builder/builder.py +++ b/src/wg_builder/builder.py @@ -16,6 +16,7 @@ random = SystemRandom() from io import BytesIO import codecs + def genPsk() -> str: return b64encode(random.randbytes(32)).decode('UTF-8') @@ -24,12 +25,13 @@ def build_network(input_file: Path = Path('/dev/stdin'), output_folder: Path = P with open(input_file) as conf: configuration = toml.load(conf) - if 'dns' in configuration: - dns = tuple(IPAddress(address) for address in configuration['dns']) - else: - dns = tuple() + dns = ( + Maybe.of_nullable(configuration.get('dns')) + .map(lambda it: tuple(IPAddress(address) for address in it)) + .or_else_get(tuple) + ) - network_name = Maybe.of_nullable(configuration.get('network_name')).or_else(input_file.stem) + network_name = Maybe.of_nullable(configuration.get('interface_name')).or_else(input_file.stem) network = IPNetwork(configuration['netmask']) address_it = network.iter_hosts() @@ -64,11 +66,24 @@ def build_network(input_file: Path = Path('/dev/stdin'), output_folder: Path = P pre_shared_keys = {(server, peer): genPsk() for peer in peers for server in servers} + def write_route_add_script(file_obj): + text = dedent(f"""\ + #!/usr/bin/env sh + ip route add {network} dev {network_name} proto kernel scope link src {network.first} + """) + file_obj.write(text) + + def write_route_del_script(file_obj): + text = dedent(f"""\ + #!/usr/bin/env sh + ip route del {network} dev {network_name} proto kernel scope link src {network.first} + """) + file_obj.write(text) def write_server_conf(server, file_obj): text = dedent(f"""\ [Interface] - Address = {server.address} + Address = {server.address}/{network.prefixlen} ListenPort = {server.port} PrivateKey = {str(server.private_key)} """) @@ -78,7 +93,7 @@ def build_network(input_file: Path = Path('/dev/stdin'), output_folder: Path = P [Peer] PublicKey = {peer.public_key} - AllowedIPs = {peer.allowed_ip} + AllowedIPs = {peer.allowed_ip}/32 PresharedKey = {pre_shared_keys[(server, peer)]} """) file_obj.write(text) @@ -89,7 +104,7 @@ def build_network(input_file: Path = Path('/dev/stdin'), output_folder: Path = P Address = {peer.allowed_ip}/32 PrivateKey = {str(peer.private_key)} """) - dns_config = (Maybe.of(dns) + dns_config = (dns .filter(lambda it: len(it) > 0) .map(lambda it: f'DNS = {', '.join((str(addr) for addr in it))}\n') .or_else('')) @@ -101,7 +116,7 @@ def build_network(input_file: Path = Path('/dev/stdin'), output_folder: Path = P PublicKey = {server.public_key} PresharedKey = {pre_shared_keys[(server, peer)]} Endpoint = {server.endpoint}:{server.port} - AllowedIPs = {str(network)} + AllowedIPs = {str(network)} """) file_obj.write(text) @@ -117,11 +132,33 @@ def build_network(input_file: Path = Path('/dev/stdin'), output_folder: Path = P tarinfo.size = bio.tell() bio.seek(0) archive.addfile(tarinfo, bio) + with BytesIO() as bio: + writer = codecs.getwriter('utf-8')(bio) + write_route_add_script(writer) + name = f'/etc/wireguard/{network_name}-add-routes.sh' + tarinfo = TarInfo(name) + tarinfo.size = bio.tell() + tarinfo.mode = 0o544 + bio.seek(0) + archive.addfile(tarinfo, bio) + with BytesIO() as bio: + writer = codecs.getwriter('utf-8')(bio) + write_route_del_script(writer) + name = f'/etc/wireguard/{network_name}-del-routes.sh' + tarinfo = TarInfo(name) + tarinfo.size = bio.tell() + tarinfo.mode = 0o544 + bio.seek(0) + archive.addfile(tarinfo, bio) else: server_folder = output_folder / f'{server.name}' makedirs(server_folder, exist_ok=True) with (open(server_folder / f'{network_name}.conf', 'w') as f): write_server_conf(server, f) + with (open(server_folder / f'{network_name}-add-routes.sh', 'w') as f): + write_route_add_script(f) + with (open(server_folder / f'{network_name}-del-routes.sh', 'w') as f): + write_route_del_script(f) for peer in peers: if archive: @@ -139,5 +176,3 @@ def build_network(input_file: Path = Path('/dev/stdin'), output_folder: Path = P makedirs(peer_folder, exist_ok=True) with (open(peer_folder / f'{network_name}.conf', 'w') as f): write_client_conf(peer, f) - -