#!/usr/bin/env python3

import sys
import datetime
import os
import socket
import subprocess
import shutil
import json


package_info = {}
rpm_tags = ['arch', 'epoch', 'name', 'release', 'vendor', 'version']

utc_time = datetime.datetime.now().strftime("%Y-%m-%d-%s/")
if not 'VZ_CRIU_REPORT_TEST' in os.environ:
    log_dir = '/vz/abrt-criu/criu-{}/'.format(utc_time)
else:
    log_dir = '/tmp/vz/abrt-criu/criu-{}/'.format(utc_time)

vz_dump_logs = '/vz/dump/'
failed_container_ids = []

def get_rpm_info(pkg):
    try:
        evrd = subprocess.check_output(['rpm', '--queryformat', '%{ARCH} %{EPOCH} %{NAME} %{RELEASE} %{VENDOR} %{VERSION}', '-q', pkg]).decode('utf-8').split()
        arch, epoch, name, release, vendor, version = evrd
        if epoch == "(none)":
           epoch = str(0)
        rpm_info = {'arch': [arch], 'epoch': [epoch], 'name': [name], 'release': [release], 'vendor': [vendor], 'version': [version]}
        package_info.update(rpm_info)
    except Exception:
        print("Package not found")
        return 'not_found'


def get_ct_proc(container_id):
    ct_proc = {}
    try:
       print('checking ct id: {}'.format(container_id))
       it_exist = subprocess.check_call(['/usr/bin/prlctl', 'status', container_id])
    except subprocess.CalledProcessError:
       # if not 0 here, it's mean that container not exist
       print('ct id: {} removed'.format(container_id))
       pass
    except OSError:
       pass
    else:
       try:
          procps = subprocess.check_output(['/usr/bin/vzps', 'faxh', '-octid,cmd', '-E', container_id]).decode()
          for process in procps.splitlines():
              if container_id not in ct_proc:
                 ct_proc[container_id] = [process.strip().strip(container_id)]
              else:
                 ct_proc[container_id].append(process.strip().strip(container_id))
       # such exception needed if binaries not exist
       except subprocess.CalledProcessError:
          return ct_proc
       except OSError:
          # empty
          return ct_proc
    return ct_proc


def copy_logs():
    phaul_log = '/var/log/phaul.log'
    phaul_service_log = '/var/log/phaul-service.log'
    shutil.copy('/etc/os-release', os.path.join(log_dir, 'os_info'))
    if os.path.exists(phaul_log):
        shutil.copy(phaul_log, os.path.join(log_dir, 'phaul.log'))
    if os.path.exists(phaul_service_log):
        shutil.copy(phaul_service_log, os.path.join(log_dir, 'phaul-service.log'))
    if os.path.exists("/var/log/vzctl.log"):
        shutil.copy("/var/log/vzctl.log", os.path.join(log_dir, 'vzctl.log'))

    if 'CRIU_WORKDIR' in os.environ and os.path.isdir(os.environ['CRIU_WORKDIR']):
        with open(os.path.join(log_dir, "criu_workdir"), "w") as f:
            f.write(os.environ['CRIU_WORKDIR'])
        failed_logs = os.environ['CRIU_WORKDIR']
        for suf in ['Dump', 'Dump.tmp', 'Dump.restore', 'Dump.fail']:
            if failed_logs.endswith(suf):
                failed_logs = failed_logs.replace("/" + suf, "")
        logs_to_pick = failed_logs
        if failed_logs.startswith('/'):
            logs_to_pick = failed_logs[1:]
        try:
            shutil.copytree(failed_logs, os.path.join(log_dir, logs_to_pick), copy_function=os.link)
        except Exception:
            # just copy files if exception
            if os.path.exists(os.path.join(log_dir, logs_to_pick)):
                shutil.rmtree(os.path.join(log_dir, logs_to_pick))
            shutil.copytree(failed_logs, os.path.join(log_dir, logs_to_pick))
        return

    # If for some reason we haven't find CRIU_WORKDIR - just
    # collect containers with issues on a dump or restore stage
    if os.path.exists(vz_dump_logs):
        try:
            out = subprocess.check_output(["grep", '-r', 'Dumping FAILED\|Restoring FAILED', vz_dump_logs])
        except:
            # Nothing to collect
            return

        for l in out.split(b"\n"):
            # We want to pick up a folder with container logs where we found migration failure message.
            # This folder can be located either in /vz/dump/ directly or inside /vz/dump/criu_dumps
            item = l.decode('utf-8').replace(vz_dump_logs, "").split('/')[0]
            if item.startswith("criu_dumps"):
                item += "/" + failed_log.split('/')[1]
            if item not in failed_container_ids:
                failed_container_ids.append(item)
                if not os.path.exists(os.path.join(log_dir, item)):
                    # src to dest with hardlink
                    try:
                        shutil.copytree(os.path.join(vz_dump_logs, item), os.path.join(log_dir, item), copy_function=os.link)
                    except Exception:
                        # just copy files if exception
                        if os.path.exists(os.path.join(log_dir, item)):
                            shutil.rmtree(os.path.join(log_dir, item))
                        try:
                            shutil.copytree(os.path.join(vz_dump_logs, item), os.path.join(log_dir, item))
                        except:
                            pass



def main():
    if not os.path.exists(log_dir):
       os.makedirs(log_dir)
    copy_logs()
    get_rpm_info('criu')

    for item in rpm_tags:
        with open(os.path.join(log_dir, 'pkg_' + item), 'a') as the_component:
             the_component.write(' '.join(package_info[item]))

    with open(os.path.join(log_dir, 'component'), 'a') as the_component:
         the_component.write(' '.join(package_info['name']))

    with open(os.path.join(log_dir, 'kernel'), 'a') as the_component:
         the_component.write(os.uname()[2])

    with open(os.path.join(log_dir, 'hostname'), 'a') as the_component:
         the_component.write(socket.gethostname())

    try:
        with open(os.path.join(log_dir, 'runlevel'), 'a') as the_component:
            the_component.write(subprocess.check_output(['runlevel']).decode())
    except:
        pass

    with open(os.path.join(log_dir, 'reason'), 'a') as the_reason:
         the_reason.write('criu migration failed')

    with open(os.path.join(log_dir, 'package'), 'a') as the_package:
         the_package.write(('{}-{}-{}').format(''.join(package_info['name']), ''.join(package_info['version']), ''.join(package_info['release'])))

    for failed_ct in failed_container_ids:
        if os.path.exists(os.path.join(log_dir, failed_ct + '/running_processes.json')):
            with open(os.path.join(log_dir, failed_ct + '/running_processes.json'), 'a') as the_failed_ct:
                the_failed_ct.write(json.dumps(get_ct_proc(failed_ct)))

    if not 'VZ_CRIU_REPORT_TEST' in os.environ:
        os.environ['DUMP_DIR'] = log_dir
        os.system("reporter-vz-prlrep")


if __name__ == '__main__':
     main()
