client.py.in

#!/usr/bin/env python3


#
#    Raspberry Pi Backup System for Classrooms - Client Script
#    Copyright (C) 2022  Kian Kasad
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License, version 3, as
#    published by the Free Software Foundation.
#
#    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 General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
#


import argparse
import atexit
import errno
import fcntl
import json
import os
import pwd
import shutil
import stat
import subprocess
import sys
import tempfile
from fcntl import LOCK_EX, LOCK_NB, LOCK_UN
from json import JSONDecodeError
from pathlib import Path
from subprocess import CalledProcessError, DEVNULL

import requests
from requests.exceptions import HTTPError


# Default IP address of the backup server if one is not provided in the stored
# configuration data
DEFAULT_SERVER_IP = 'm4_SERVER_IP'

# Port on which the backup server is serving configuration
CONFIG_SERVER_PORT = m4_HTTPD_PORT

# Port on which the backup server is listening for SSH connections
SSHD_PORT = m4_SSHD_PORT

# Location of drop-in configuration file for backup.timer
TIMER_DROPIN_PATH = '/etc/systemd/system/backup.timer.d/00-times.conf'

# Default location for the backup client's lock file. If $RUNTIME_DIRECTORY
# is set, the file will be placed in that directory instead.
DEFAULT_LOCK_FILE_PATH = '/var/run/backup_client.lock'

# Command to use instead of 'ssh'
rsh = 'ssh -q -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null'
SSH_KEY = '/usr/local/share/backup_client/ssh_key'

# Environment variables to run Borg with
BORG_ENV = {
    'BORG_UNKNOWN_UNENCRYPTED_REPO_ACCESS_IS_OK': 'yes',
    'BORG_RELOCATED_REPO_ACCESS_IS_OK': 'yes',
}

# Empty configuration structure.
# Must contain all the keys that might be accessed but need not contain any
# values.
DEFAULT_CONFIG = {
    'epoch': 0,
    'server': {
        'host': DEFAULT_SERVER_IP,
        'httpd_port': CONFIG_SERVER_PORT,
        'sshd_port': SSHD_PORT,
    },
    'updates': [
        {
            'epoch': 0,
            'script': 'updates/00-no_op.sh',
        }
    ],
    'archive_name_format': '{now}',
    'backup_times': [
        '@before:shutdown.target'
    ],
    'backup_time_randomized_delay': '2min',
    'backup_paths': ['~/Desktop'],
    'backup_user': 'pi',
}

# Lock file handle
lockfile = None

# Temporary SSH key path
tmp_ssh_key = None


def die(*args, **kwargs):
    if args:
        error(*args, **kwargs)
    print('Exiting...')
    exit(1)


def log(*args, **kwargs):
    print('\x1b[1;34mInfo:\x1b[m ', end='')
    print(*args, **kwargs)


def warn(*args, **kwargs):
    print('\x1b[1;33mWarning:\x1b[m ', end='')
    print(*args, **kwargs)


def error(*args, **kwargs):
    print("\x1b[1;31mError:\x1b[m ", end='')
    print(*args, **kwargs)


def fill_struct_with_defaults(source, default):
    for key, value in default.items():
        if key not in source:
            source[key] = value
        elif isinstance(source[key], dict):
            fill_struct_with_defaults(source[key], value)
        elif isinstance(source[key], str) and not source[key]:
            source[key] = value


def time_is_valid(type, time):
    VALID_TYPES = ['timespan', 'calendar']
    if type not in VALID_TYPES:
        raise ValueError('Invalid time type')
    cmd = ['/usr/bin/systemd-analyze', type, time]
    proc = subprocess.run(cmd, stdout=DEVNULL, stderr=DEVNULL)
    return proc.returncode == 0


def apply_config(config, prev_config, skip=0):
    log('Applying new configuration...')

    # Apply updates from updates list
    current_update_epoch = max(max(
        [update['epoch'] for update in prev_config['updates']]
    ), skip)
    latest_update_epoch = max(
        [update['epoch'] for update in config['updates']]
    )
    for i in range(current_update_epoch + 1, latest_update_epoch + 1):
        try:
            update = [update for update in config['updates']
                      if update['epoch'] == i][0]
            apply_update(config['server'], i, update['script'])
        except IndexError:
            pass

    # Set times for backup.timer
    # See systemd.timer(5) and systemd.time(7) for details
    enable_on_shutdown = False
    try:
        with open(TIMER_DROPIN_PATH, 'w') as dropin:
            print('[Timer]', file=dropin)
            # Print randomized delay setting
            time = config['backup_time_randomized_delay']
            if time_is_valid('timespan', time):
                print('RandomizedDelaySec=', time, sep='', file=dropin)
            else:
                warn(f"Time span '{time}' is invalid. Ignoring it...")

            # Print times
            print('OnCalendar=', file=dropin)
            for time in config['backup_times']:
                if time == '@shutdown':
                    enable_on_shutdown = True
                elif time_is_valid('calendar', time):
                    print('OnCalendar=', time, sep='', file=dropin)
                else:
                    warn(f"Calendar time '{time}' is invalid. Ignoring it...")
    except OSError as err:
        die('Failed to save trigger times to ',
            TIMER_DROPIN_PATH, ': ', err, sep='')

    # Enable/disable backup-on-shutdown.service
    cmd = [
        'systemctl', '--no-ask-password',
        'enable' if enable_on_shutdown else 'disable',
        '--now',
        'backup-on-shutdown.service'
    ]
    try:
        subprocess.run(cmd, check=True, stdout=DEVNULL, stderr=DEVNULL)
    except CalledProcessError as err:
        warn('Failed to enable backup on shutdown service:', err)

    # Reload systemd
    try:
        cmd = ['systemctl', '--no-ask-password', 'daemon-reload']
        subprocess.run(cmd, check=True, stdout=DEVNULL, stderr=DEVNULL)
    except CalledProcessError as err:
        warn('Failed to reload systemd:', err)
        warn('Continuing anyways...')


def apply_update(server, epoch, script_name):
    # Log and flush buffers
    log('Attempting to apply update #', epoch, '...', sep='')
    sys.stdout.flush()
    sys.stderr.flush()

    uri = f"http://{server['host']}:{server['httpd_port']}/{script_name}"
    try:
        response = requests.get(uri)
        script = response.text
        status = subprocess.run(
            '/bin/bash', input=script, text=True, check=True)
    except CalledProcessError as err:
        die(f"Update #{epoch}'s script returned non-zero exit status {err.returncode}")
    except ConnectionError as err:
        die(f'Failed to download script for update #{epoch}:', err)
    except HTTPError as err:
        die('HTTP error:', err)

    log(f'Applied update #{epoch}.')
    log('Re-executing backup client...')
    reexec_script('--resume-after-update', epoch)


# Re-execute this client script
def reexec_script(*args):
    global lockfile

    # Release lockfile
    try:
        fcntl.flock(lockfile, LOCK_UN | LOCK_NB)
    except OSError:
        warn('Failed to release lock on', lockfile.name)
    lockfile.close()

    delete_ssh_key()

    # Flush output buffers
    sys.stdout.flush()
    sys.stderr.flush()

    # Execute script
    progargs = [sys.argv[0]] + [str(arg) for arg in args]
    try:
        os.execvp(sys.argv[0], progargs)
    except OSError as err:
        error('Failed to re-execute backup client:', err)
        die('Manually clean up system before running backup client again.')


def store_config(data, path):
    log('Storing updated configuration file...')
    try:
        with open(path, 'w') as file:
            file.write(data)
    except OSError as err:
        warn('Failed to save updated configuration:', err)


def do_backup(config, is_retry=False):
    global rsh, BORG_ENV

    log('Running Borg backup...')

    server_ip = config['server']['host']
    sshd_port = config['server']['sshd_port']

    repo_uri = f'ssh://backup@{server_ip}:{sshd_port}/~/repos/' + '{hostname}'
    archive_path = repo_uri + '::' + config['archive_name_format']

    cmd = [
        '/usr/bin/borg',
        '--rsh', rsh,
        'create',
        '--log-json',
        '--json',
        archive_path,
        *[os.path.expanduser(path) for path in config['backup_paths']]
    ]

    # Run Borg
    proc = subprocess.run(cmd, env=BORG_ENV, capture_output=True, text=True)

    # Process log messages
    repo_doesnt_exist = False
    errors = []
    warnings = []
    for line in proc.stderr.splitlines():
        try:
            msg = json.loads(line)
        except JSONDecodeError:
            continue  # ignore non-JSON output
        if msg['levelname'] == 'ERROR':
            if msg.get('msgid', '') == 'Repository.DoesNotExist':
                repo_doesnt_exist = True
            errors.append(msg)
        elif msg['levelname'] == 'WARNING':
            warnings.append(msg)

    if proc.returncode == 0:
        # Succeeded
        log('Borg successfully created backup.')
    elif proc.returncode == 1:
        # Succeeded with warnings
        warn('Borg successfully created backup, but produced',
             len(warnings), 'warnings (listed below).')
        for msg in warnings:
            warn('borg:', msg['message'])
    elif proc.returncode == 2:
        # Failed
        if repo_doesnt_exist and not is_retry:
            warn('Backup repository does not exist. Creating new repository...')
            create_repo(repo_uri)
            do_backup(config, True)
        else:
            error('Borg failed to create backup and produced',
                  len(errors), 'errors and', len(warnings), 'warnings (listed below).')
            for msg in errors:
                error('borg:', msg['message'])
            for msg in warnings:
                warn('borg:', msg['message'])
            die()


def create_repo(uri):
    global rsh, BORG_ENV
    cmd = [
        '/usr/bin/borg',
        '--rsh', rsh,
        'init',
        '--log-json',
        '--encryption', 'none',
        uri
    ]

    # Run Borg
    proc = subprocess.run(cmd, env=BORG_ENV, capture_output=True, text=True)

    # Process log messages
    errors = []
    warnings = []
    for line in proc.stderr.splitlines():
        try:
            msg = json.loads(line)
        except JSONDecodeError:
            continue  # ignore non-JSON output
        if msg['levelname'] == 'ERROR':
            errors.append(msg)
        elif msg['levelname'] == 'WARNING':
            warnings.append(msg)

    # Check return code
    if proc.returncode == 0:
        # Success
        log('Successfully created repository at', uri)
    elif proc.returncode == 1:
        warn('Successfully created repository at', uri,
             'but Borg produced', len(warnings), warnings)
        for msg in warnings:
            warn('borg:', msg['message'])
    elif proc.returncode == 2:
        # Failed
        error('Borg failed to create repository and produced',
              len(warnings), 'warnings and', len(errors), 'errors (listed below).')
        for msg in warnings:
            error('borg:', msg['message'])
        for msg in errors:
            warn('borg:', msg['message'])
        die()


# Make a user-readable copy of the SSH key.
# Stores the path to the copy in `tmp_ssh_key'.
def copy_ssh_key(uid, gid):
    global tmp_ssh_key
    try:
        # Create temporary key file
        (fd, tmp_ssh_key) = tempfile.mkstemp()
        os.chown(fd, uid, gid)
        os.chmod(fd, stat.S_IRUSR)

        # Delete temporary key file when program exits
        atexit.register(delete_ssh_key)

        # Copy contents from original key
        with open(SSH_KEY, 'rb') as src, open(fd, 'wb') as dst:
            shutil.copyfileobj(src, dst)
    except Exception as err:
        die('Failed to make copy of SSH key:', err)
    return tmp_ssh_key


def delete_ssh_key():
    global tmp_ssh_key
    if tmp_ssh_key:
        try:
            os.remove(tmp_ssh_key)
        except OSError as err:
            warn('Failed to delete temprary SSH key:', err)


def get_drop_ids(user):
    try:
        if isinstance(user, int):
            uid = user
            gid = pwd.getpwuid(uid).pw_gid
        else:
            pwinfo = pwd.getpwnam(user)
            uid = pwinfo.pw_uid
            gid = pwinfo.pw_gid
    except KeyError as err:
        die('Failed to get UID/GID of target user:', err)
    return (uid, gid)


def drop_privileges_to(uid, gid):
    try:
        # Order is important
        os.environ.clear()
        os.setgroups([])
        os.setgid(gid)
        os.setuid(uid)
        os.umask(0o022)
        os.chdir(Path.home())
    except Exception as err:
        die('Failed to drop privileges:', err)


# Attempt to acquire lock file
def acquire_lock_file():
    global DEFAULT_LOCK_FILE_PATH, lockfile
    try:
        lockfile_path = Path(os.getenv('RUNTIME_DIRECTORY').split(':')[0])
        lockfile_path /= 'backup_client.lock'
    except:
        warn('Environment variable RUNTIME_DIRECTORY is not set. Falling back to',
             DEFAULT_LOCK_FILE_PATH, 'as lock file.')
        lockfile_path = DEFAULT_LOCK_FILE_PATH
    try:
        lockfile = open(lockfile_path, 'w')
    except OSError as err:
        if lockfile:
            lockfile.close()
        die('Failed to open lock file:', err)
    try:
        fcntl.flock(lockfile, LOCK_EX | LOCK_NB)
    except OSError as err:
        if err.errno == errno.EWOULDBLOCK or err.errno == errno.EACCES:
            die(f'Cannot acquire lock on {DEFAULT_LOCK_FILE_PATH}. Is the backup client is already running?')
        else:
            die('Failed to acquire lock on ', err, '.', sep='')


def stored_config_path():
    try:
        state_dir = Path(os.getenv('STATE_DIRECTORY').split(':')[0])
    except:
        die("Environment variable STATE_DIRECTORY is not set")
    return state_dir / 'config.json'


# Attempt to load stored configuration file
def load_stored_config():
    global DEFAULT_CONFIG
    try:
        with open(stored_config_path(), 'r') as file:
            stored_config = json.load(file)
            fill_struct_with_defaults(stored_config, DEFAULT_CONFIG)
    except FileNotFoundError:
        warn('No stored configuration file found.')
        warn('Continuing with default settings...')
        stored_config = DEFAULT_CONFIG
    except ValueError:
        warn('Stored configuration contains invalid JSON.')
        warn('Continuing with defaults instead...')
        stored_config = DEFAULT_CONFIG
    except Exception as err:
        warn('Unable to open stored configuration file:', err)
        warn('Continuing with defaults instead...')
        stored_config = DEFAULT_CONFIG
    return stored_config


# Parse command-line arguments
def parse_command_line_args():
    ap = argparse.ArgumentParser()
    ap.add_argument('-u', '--resume-after-update',
                    help='Skips application of updates prior to the epoch given',
                    type=int,
                    default=0)
    return ap.parse_args()


def fetch_configuration_from(host, port):
    global DEFAULT_CONFIG
    uri = f'http://{host}:{port}/config.json'
    log(f'Fetching configuration from {uri}...')
    try:
        response = requests.get(uri)
        response.raise_for_status()
        config_json = response.text
        config = response.json()
        fill_struct_with_defaults(config, DEFAULT_CONFIG)
        return (config, config_json)
    except HTTPError as err:
        error('HTTP error:', err)
    except ValueError as err:
        error('Invalid JSON response from configuration server:', err)
    except Exception as err:
        error('Failed to fetch configuration:', err)
    return None


def main():
    global DEFAULT_CONFIG, DEFAULT_SERVER_IP, CONFIG_SERVER_PORT, rsh

    acquire_lock_file()

    args = parse_command_line_args()
    skip_updates_through = args.resume_after_update

    stored_config = load_stored_config()

    # Fetch configuration
    resp = fetch_configuration_from(stored_config['server']['host'],
                                    stored_config['server']['httpd_port'])
    if resp is None:
        warn('Failed to fetch configuration.',
             'Retrying with default server parameters...')
        resp = fetch_configuration_from(DEFAULT_SERVER_IP, CONFIG_SERVER_PORT)
    if resp is None:
        die('Failed to fetch configuration from server.')
    else:
        config, config_json = resp

    # If fetched config is newer, apply and store it
    if config['epoch'] > stored_config['epoch']:
        apply_config(config, stored_config, skip=skip_updates_through)
        store_config(config_json, stored_config_path())

    # Drop privileges
    uid, gid = get_drop_ids(config['backup_user'])
    tmpkey = copy_ssh_key(uid, gid)
    rsh += ' -i ' + tmpkey
    drop_privileges_to(uid, gid)

    do_backup(config)


if __name__ == '__main__':
    main()