FreedomBox/actions/wireguard

#!/usr/bin/python3
#
# This file is part of FreedomBox.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
"""
Configuration helper for WireGuard.
"""

import argparse
import json
import subprocess

PUBLIC_KEY_HELP = 'Public key for the client'

SERVER_INTERFACE = 'wg0'


def parse_arguments():
    """Return parsed command line arguments as dictionary."""
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest='subcommand', help='Sub command')

    subparsers.add_parser('setup', help='Setup WireGuard')
    subparsers.add_parser('get-info', help='Get server and clients info')

    add_client = subparsers.add_parser('add-client', help='Add a client')
    add_client.add_argument('publickey', help=PUBLIC_KEY_HELP)

    remove_client = subparsers.add_parser('remove-client',
                                          help='Remove a client')
    remove_client.add_argument('publickey', help=PUBLIC_KEY_HELP)

    add_server = subparsers.add_parser('add-server', help='Add a server')
    add_server.add_argument('--endpoint', required=True,
                            help='Server endpoint')
    add_server.add_argument('--client-ip', required=True,
                            help='Client IP address provided by server')
    add_server.add_argument('--public-key', required=True,
                            help='Public key of the server')
    add_server.add_argument('--pre-shared-key', help='Pre-shared key')
    add_server.add_argument(
        '--all-outgoing', action='store_true',
        help='Use this connection to send all outgoing traffic')

    subparsers.required = True
    return parser.parse_args()


def subcommand_setup(_):
    """Setup WireGuard."""
    subprocess.run(
        ['ip', 'link', 'add', 'dev', SERVER_INTERFACE, 'type', 'wireguard'],
        check=True)
    subprocess.run(
        ['wg', 'set', SERVER_INTERFACE, 'listen-port', '51820'], check=True)
    # TODO: generate key pair


def subcommand_get_info(_):
    """Get server and clients info."""
    output = subprocess.check_output(
        ['wg', 'show', SERVER_INTERFACE, 'dump']).decode().strip()
    lines = output.split('\n')
    server_data = lines.pop(0).split()
    server = {
        'private_key': server_data[0],
        'public_key': server_data[1],
        'listen_port': server_data[2],
        'fwmark': server_data[3],
    }

    clients = []
    for client_line in lines:
        client_data = client_line.split()
        client_info = {
            'public_key': client_data[0],
            'preshared_key': client_data[1],
            'endpoint': client_data[2],
            'allowed_ips': client_data[3],
            'latest_handshake': client_data[4],
            'transfer_rx': client_data[5],
            'transfer_tx': client_data[6],
            'persistent_keepalive': client_data[7],
        }
        clients.append(client_info)

    # TODO: Add servers info from other interfaces.
    info = {
        'server': server,
        'clients': clients,
    }
    print(json.dumps(info))


def subcommand_add_client(arguments):
    """Add a client."""
    subprocess.run(
        ['wg', 'set', SERVER_INTERFACE, 'peer', arguments.publickey],
        check=True)


def subcommand_remove_client(arguments):
    """Remove a client."""
    subprocess.run(
        ['wg', 'set', SERVER_INTERFACE, 'peer', arguments.publickey, 'remove'],
        check=True)


def subcommand_add_server(arguments):
    """Add a server."""
    output = subprocess.check_output(
        ['wg', 'show', 'interfaces']).decode().strip()
    interfaces = output.split()
    interface_num = 1
    for interface in interfaces:
        new_interface_name = 'wg' + str(interface_num)
        if interface == new_interface_name:
            interface_num += 1
        else:
            break

    subprocess.run(
        ['ip', 'link', 'add', 'dev', new_interface_name, 'type', 'wireguard'],
        check=True)

    args = ['wg', 'set', interface, 'peer', arguments.public_key]
    if arguments.pre_shared_key:
        args += ['preshared-key', arguments.pre_shared_key]

    args += ['endpoint', arguments.endpoint]
    subprocess.run(args, check=True)


def main():
    """Parse arguments and perform all duties."""
    arguments = parse_arguments()

    subcommand = arguments.subcommand.replace('-', '_')
    subcommand_method = globals()['subcommand_' + subcommand]
    subcommand_method(arguments)


if __name__ == '__main__':
    main()