fetchlogs

#!/usr/bin/env python3

import argparse
import os
import shutil
import subprocess
import sys
import zipfile


p = subprocess.Popen("docker compose version > /dev/null 2>&1", shell=True)
p.wait()
if p.returncode != 0:
    CNAME_SEP = '_'
    DC_COMMAND = 'docker-compose -p strato --log-level=ERROR'
    print("Docker Compose V1. Using '_' as a container name separator")
else:
    CNAME_SEP = '-'
    DC_COMMAND = 'docker compose -p strato'
    print("Docker Compose V2. Using '-' as a container name separator")
    
def fetch_logs():
    containers_with_logfiles = {
        'bloc': '/logs',
        'strato': '/var/lib/strato/logs'
    }
    if not os.path.isfile('docker-compose.yml'):
        sys.stderr.write('ERROR: ./docker-compose.yml does not exist. Try `./strato --compose` to download one')
        sys.exit(6)
    try:
        compose_ps_output = subprocess.check_output(
            "%s ps --services" % DC_COMMAND,
            shell=True
        ).decode('utf-8').split('\n')
        running_containers = ['strato%(sep)s%(service)s%(sep)s1' % {'service': service_name, 'sep': CNAME_SEP} for service_name in compose_ps_output if service_name]
    except Exception as e:
        sys.stderr.write('ERROR: Unable to fetch the currently running containers:\n%s' % e)
        sys.exit(7)

    for container_name in running_containers:
        container_name_split = container_name.split(CNAME_SEP)
        if len(container_name_split) > 3:
            container_name_split = [container_name_split[0], CNAME_SEP.join(str(item) for item in container_name_split[1:-1]) , container_name_split[-1]]
        elif len(container_name_split) != 3:
            sys.stderr.write('ERROR: Unexpected container name: %s' % container_name)
            sys.exit(8)
        container_subname = container_name_split[1]
        if container_subname in containers_with_logfiles.keys():
            current_container_logs_dir = os.path.join(temp_dir_path, container_subname)
            os.makedirs(current_container_logs_dir)
            docker_cp_command = 'docker cp %s:%s %s' % (
                container_name,
                os.path.join(containers_with_logfiles[container_subname], '.'),
                current_container_logs_dir
            )
            try:
                subprocess.check_output(docker_cp_command, shell=True)
            except Exception as e:
                error_message = 'ERROR: Unable to fetch logs from %s container with error:\n%s' % (container_name, e)
                sys.stderr.write('%s\n' % error_message)
                create_text_file(os.path.join(current_container_logs_dir, '%s_error.txt' % container_subname),
                                 error_message)
        else:
            # Copying logfile returned with `"docker inspect --format='{{.LogPath}}' %s" % container_name` does not work
            # since: 1. it returns non-existing log filenames in some cases 2. sudo is mandatory to access file directly
            try:
                subprocess.check_output(
                    'docker logs -t %s > %s 2>&1' % (container_name, os.path.join(temp_dir_path, container_subname)),
                    shell=True
                )
            except Exception as e:
                error_message = 'ERROR: Unable to fetch logs of %s container with error:\n%s' % (container_name, e)
                sys.stderr.write('%s\n' % error_message)
                create_text_file(os.path.join(temp_dir_path, '%s_error.txt' % container_subname), error_message)
                deferred_errors.append(error_message)

    if args.db_dump:
        postgres_container_name = [cname for cname in running_containers if '%(sep)spostgres%(sep)s' % {'sep': CNAME_SEP} in cname][0]
        sys.stdout.write('Creating Posgtres DBs dump - this may take up to several minutes...\n')
        pgdump_filename = 'pgdump'
        try:
            subprocess.check_output(
                'docker exec -i %s pg_dumpall -U postgres -f %s && '
                'docker cp %s:/tmp/pgdump %s && '
                'docker exec -i %s rm -rf /tmp/pgdump' % (
                    postgres_container_name,
                    os.path.join('/tmp', pgdump_filename),
                    postgres_container_name,
                    temp_dir_path,
                    postgres_container_name
                ),
                shell=True
            )
        except Exception as e:
            error_message = 'ERROR: unable to process the database dump with error:\n%s' % e
            sys.stderr.write('%s\n' % error_message)
            create_text_file(os.path.join(temp_dir_path, '%s_error.txt' % pgdump_filename), error_message)
            deferred_errors.append(error_message)

    if not args.as_dir:
        try:
            sys.stdout.write('Zipping data - almost done...\n')
            if args.force:
                try:
                    os.remove(output_zip_path)
                except OSError:
                    pass
            zip_logs(output_zip_name)
        except Exception as e:
            sys.stderr.write(
                'ERROR: Unable to archive the directory %s into zip archive. Check the logs in directory: %s\nError:\n%s' % (
                    archive_output_dir,
                    temp_dir_path,
                    e
                ))
            sys.exit(9)
        try:
            shutil.rmtree(temp_dir_path)
        except Exception as e:
            sys.stderr.write(
                'ERROR: Unable to clean the temporary directory (%s) with the error:\n%s\n' % (temp_dir_path, e))


def create_text_file(file_path, text):
    try:
        with open(file_path, "w") as f:
            f.write(text)
    except Exception as e:
        error_message = 'SCRIPT ERROR: unable to create file %s with an error: %s' % (file_path, e)
        sys.stderr.write('%s\n' % error_message)
        deferred_errors.append(error_message)


def zip_logs(output_file_name):
    work_dir = os.getcwd()
    os.chdir(archive_output_dir)
    with zipfile.ZipFile(output_file_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(temp_dir_name):
            for file_to_zip in files:
                zipf.write(os.path.join(root, file_to_zip))
    os.chdir(work_dir)


if __name__ == "__main__":
    deferred_errors = []
    parser = argparse.ArgumentParser(
        description='Fetch all STRATO logs and optionally the database dump into single zip file',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        '-o', '--output-dir',
        dest='path',
        default='./',
        help='Output directory path',
    )
    parser.add_argument(
        '--db-dump',
        action='store_true',
        help='Add the Postgres database dump. WARNING: the data stored in databases may be sensitive; '
             'the dump process may lock the databases for up to several minutes; the dump file may be large',
    )
    parser.add_argument(
        '-f', '--force',
        action='store_true',
        help='Force rewrite to a previously exported log files and directories',
    )
    parser.add_argument(
        '--as-dir',
        action='store_true',
        help='Leave the logs as a directory instead of a zip file',
    )

    args = parser.parse_args()
    archive_output_dir = args.path
    if not os.path.isdir(archive_output_dir):
        try:
            os.makedirs(archive_output_dir)
        except Exception as e:
            sys.stderr.write('ERROR: Unable to create directory %s:\n%s' % (archive_output_dir, e))
            sys.exit(2)
    if not os.access(archive_output_dir, os.W_OK):
        sys.stderr.write(
            'ERROR: No write access to: "%s". Change directory permissions or provide another path as a script argument')
        sys.exit(3)

    temp_dir_name = 'strato_logs'
    temp_dir_path = os.path.join(archive_output_dir, temp_dir_name)

    output_zip_name = 'strato_logs.zip'
    output_zip_path = os.path.join(archive_output_dir, output_zip_name)
    if args.force:
        try:
            shutil.rmtree(temp_dir_path)
        except OSError:
            pass
        os.makedirs(temp_dir_path)
    else:
        if os.path.isfile(output_zip_path):
            error_message = 'Output zip file already exists (%s)' % os.path.abspath(output_zip_path)
            overwrite_input = input('%s, remove? (y/n): ' % error_message)[0].lower()
            if overwrite_input == 'y':
                os.remove(output_zip_path)
            else:
                sys.stderr.write(error_message)
                sys.exit(4)
        try:
            os.makedirs(temp_dir_path)
        except OSError as e:
            error_message = 'Temporary directory exists (%s)' % os.path.abspath(temp_dir_path)
            overwrite_input = input('%s, overwrite? (y/n): ' % error_message)[0].lower()
            if overwrite_input == 'y':
                shutil.rmtree(temp_dir_path)
                os.makedirs(temp_dir_path)
            else:
                sys.stderr.write(error_message)
                sys.exit(5)

    sys.stdout.write('Fetching STRATO logs - this may take up to a minute...\n')
    fetch_logs()

    if deferred_errors:
        sys.stderr.write('ERROR: Some logs failed to fetch with errors:\n%s' % '\n'.join(deferred_errors))
        sys.exit(100)
    else:
        sys.stdout.write('Done.\n')
        sys.exit(0)