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 os
import pathlib
import subprocess

PUBLIC_KEY_HELP = 'Public key for the client'

SERVER_INTERFACE = 'wg0'

KEY_FOLDER = pathlib.Path('/var/lib/freedombox/wireguard')
PRIVATE_KEY_PATH = KEY_FOLDER / 'privatekey'
PUBLIC_KEY_PATH = KEY_FOLDER / 'publickey'


class InterfaceNotFoundError(Exception):
    """Exception raised when no matching interface is found."""


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 info for each configured interface')

    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')

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

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

    subparsers.required = True
    return parser.parse_args()


def _generate_key_pair():
    """Generate private/public key pair."""
    private_key = subprocess.check_output(['wg', 'genkey'])
    public_key = subprocess.check_output(['wg', 'pubkey'], input=private_key)
    KEY_FOLDER.mkdir(parents=True, exist_ok=True)
    with PUBLIC_KEY_PATH.open(mode='wb') as public_key_file:
        public_key_file.write(public_key)

    old_umask = os.umask(0o077)
    try:
        with PRIVATE_KEY_PATH.open(mode='wb') as private_key_file:
            private_key_file.write(private_key)

    finally:
        os.umask(old_umask)


def subcommand_setup(_):
    """Setup WireGuard."""
    # Create interface.
    try:
        subprocess.run(['ip', 'link', 'show', SERVER_INTERFACE],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                       check=True)
    except subprocess.CalledProcessError:
        subprocess.run(
            ['ip', 'link', 'add', 'dev', SERVER_INTERFACE, 'type',
             'wireguard'], check=True)

    if not (PUBLIC_KEY_PATH.exists() and PRIVATE_KEY_PATH.exists()):
        _generate_key_pair()

    # Configure interface.
    subprocess.run(
        ['wg', 'set', SERVER_INTERFACE, 'listen-port', '51820', 'private-key',
         str(PRIVATE_KEY_PATH)], check=True)


def _get_info():
    """Return info for each configured interface."""
    output = subprocess.check_output(
        ['wg', 'show', 'all', 'dump']).decode().strip()
    lines = output.split('\n')
    interfaces = {}
    for line in lines:
        fields = line.split()
        interface_name = fields[0]
        if interface_name in interfaces:
            peer = {
                'public_key': fields[1],
                'preshared_key': fields[2],
                'endpoint': fields[3],
                'allowed_ips': fields[4],
                'latest_handshake': fields[5],
                'transfer_rx': fields[6],
                'transfer_tx': fields[7],
                'persistent_keepalive': fields[8],
            }
            interfaces[interface_name]['peers'].append(peer)

        else:
            interfaces[interface_name] = {
                'interface_name': interface_name,
                'private_key': fields[1],
                'public_key': fields[2],
                'listen_port': fields[3],
                'fwmark': fields[4],
                'peers': [],
            }

    return interfaces


def subcommand_get_info(_):
    """Print info for each configured interface."""
    print(json.dumps(_get_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 _find_next_interface():
    """Find next unused wireguard interface name."""
    output = subprocess.check_output(
        ['wg', 'show', 'interfaces']).decode().strip()
    interfaces = output.split()
    interface_num = 1
    new_interface_name = 'wg1'
    while new_interface_name in interfaces:
        interface_num += 1
        new_interface_name = 'wg' + str(interface_num)

    return new_interface_name


def _create_connection(name, interface, client_ip):
    """Create a NetworkManager connection."""
    subprocess.run(['nmcli', 'con', 'add',
                    'con-name', name,
                    'ifname', interface,
                    'type', 'wireguard'], check=True)

    subprocess.run(['nmcli', 'con', 'modify', name,
                    'connection.autoconnect', 'TRUE'], check=True)

    subprocess.run(['nmcli', 'con', 'modify', name,
                    'connection.zone', 'internal'], check=True)

    subprocess.run(['nmcli', 'con', 'modify', name,
                    'ipv4.method', 'manual',
                    'ipv4.addresses', client_ip + '/24'], check=True)

    with PRIVATE_KEY_PATH.open() as private_key_file:
        private_key = private_key_file.read().strip()

    subprocess.run(['nmcli', 'con', 'modify', name,
                    'wireguard.private-key', private_key], check=True)


def subcommand_add_server(arguments):
    """Add a server."""
    new_interface_name = _find_next_interface()

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

    connection_name = 'WireGuard-' + new_interface_name
    _create_connection(connection_name, new_interface_name,
                       arguments.client_ip)

    # XXX: Peer is lost after connection is activated.
    args = ['wg', 'set', new_interface_name, '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 subcommand_modify_server(arguments):
    """Modify a server."""
    interfaces = _get_info()
    interfaces.pop(SERVER_INTERFACE)
    interface_to_modify = None
    for interface in interfaces.values():
        if interface['peers']:
            peer = interface['peers'][0]
            if peer['public_key'] == arguments.public_key:
                interface_to_modify = interface['interface_name']

    if interface_to_modify:
        args = ['wg', 'set', interface_to_modify, '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)

        subprocess.run(['nmcli', 'con', 'modify',
                        'WireGuard-' + interface_to_modify,
                        'ipv4.method', 'manual',
                        'ipv4.addresses', arguments.client_ip + '/24'],
                       check=True)

    else:
        raise InterfaceNotFoundError(
            'Interface with peer %s not found' % arguments.publickey)


def subcommand_remove_server(arguments):
    """Remove a server."""
    interfaces = _get_info()
    interfaces.pop(SERVER_INTERFACE)
    interface_to_remove = None
    for interface in interfaces.values():
        if interface['peers']:
            peer = interface['peers'][0]
            if peer['public_key'] == arguments.publickey:
                interface_to_remove = interface['interface_name']

    if interface_to_remove:
        subprocess.run(
            ['nmcli', 'con', 'delete', 'WireGuard-' + interface_to_remove],
            check=True)
        subprocess.run(['ip', 'link', 'delete', interface_to_remove],
                       check=True)

    else:
        raise InterfaceNotFoundError(
            'Interface with peer %s not found' % arguments.publickey)


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()