#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# vz-problem-report: Generate problem report for Virtuozzo.
#
# Copyright (c) 2020 Virtuozzo International GmbH. All rights reserved.
#
# Our contact details: Virtuozzo International GmbH, Vordergasse 59, 8200
# Schaffhausen, Switzerland.

# standart python libs
import re, os, sys, shutil, base64, logging, locale
import tarfile, argparse, tempfile, subprocess

# additional python libs, must be present
import urllib3, requests

# special python libs, can be absent
try:
    import winrm
except ImportError:
    pass

try:
    import libvirt
except ImportError:
    pass

report_arch = 'report.tar.gz'
report_server = 'ps6.report.virtuozzo.com'
report_xml_templet = '''<VirtuozzoProblemReport>
  <Type>17</Type>
  <ProductName>Docker libvirt {}</ProductName>
  <ProblemDescription>{}</ProblemDescription>
  <ClientVersion>1.03</ClientVersion>
  <ContactInfo>
    <Name>{}</Name>
    <EMail>{}</EMail>
  </ContactInfo>
  <ClientProxyInfo>{}</ClientProxyInfo>
  <AppConfig>{}</AppConfig>
</VirtuozzoProblemReport>
'''

########################
# args option handling #
########################

def parse_args():
    parser = argparse.ArgumentParser(description = 'Generate and send problem report to Virtuozzo.'
                                     ' Also you can save archive with report in file or dir.')

    parser.add_argument('address', help = 'an address for destination host')
    parser.add_argument('-v', '--verbose', help = 'print verbose info', action='count', default=0)
    parser.add_argument('--proxy', help = 'proxy in format user:password@proxyhost:port')
    parser.add_argument('--name', help = 'your name for saving it in report', default = '')
    parser.add_argument('--email', help = 'your e-mail for saving it in report', default = '')
    parser.add_argument('--description', help = 'desription of your problem', metavar = 'DESC', default = '')
    parser.add_argument('--reportfile', help = 'don\'t send report, only save report in PATH', metavar = 'PATH')
    parser.add_argument('--version', version='%(prog)s 1.03', action = 'version')

    return parser.parse_args()

#######################
# cmd output handling #
#######################

def cmd_out(cmd):
    try:
        out = subprocess.check_output(cmd.split())
    except Exception:
        logging.error('Execution of command failed')
        return 'Command "{}" failed'.format(cmd)

    return out.decode(locale.getpreferredencoding(), errors = 'ignore')

########################
# certificate handling #
########################

def get_cert_dir():
    def check_dir(cert_dir):
        return (os.path.isfile(os.path.join(cert_dir, 'cert.pem')) and
                os.path.isfile(os.path.join(cert_dir, 'key.pem')) and
                os.path.isfile(os.path.join(cert_dir, 'ca.pem')))

    cert_dir = os.getenv('DOCKER_CERT_PATH')
    if not cert_dir is None and check_dir(cert_dir):
        return cert_dir

    if os.name != 'posix' or os.geteuid() != 0:
        cert_dir = os.path.expanduser('~/.docker')
    else:
        cert_dir = '/etc/pki/docker'

    if check_dir(cert_dir):
        return cert_dir

    return '/etc/pki/libvirt/docker'

#############################
# winrm connection handling #
#############################

def open_winrm(address):
    cert_dir = get_cert_dir()

    try:
        protocol = winrm.Protocol('https://{0}:5986/wsman'.format(address),
            transport='certificate',
            cert_pem = os.path.join(cert_dir, 'cert.pem'),
            cert_key_pem = os.path.join(cert_dir, 'key.pem'),
            ca_trust_path = os.path.join(cert_dir, 'ca.pem'))

        shell_id = protocol.open_shell()
    except Exception:
        logging.error('Connecting winrm to "%s" failed', address)
        return None

    return (protocol, shell_id)

def close_winrm(winrm):
    protocol, shell_id = winrm
    protocol.close_shell(shell_id)

#############################
# remote execution handling #
#############################

def remote_cmd(winrm, cmd, args):
    protocol, shell_id = winrm

    command_id = protocol.run_command(shell_id, cmd, args)
    out, code = protocol.get_command_output(shell_id, command_id)[::2]
    protocol.cleanup_command(shell_id, command_id)

    if code != 0:
        if cmd == 'powershell':
            return None
        else:
            logging.error('Execution of remote command failed')
            return 'Remote command "{}" failed\n'.format(cmd)

    return out.decode('cp437', errors = 'ignore')

def remote_ps(winrm, ps):
    encoded_ps = base64.b64encode(ps.encode('utf_16_le')).decode('ascii')
    logging.debug('PS: %s', ps)
    logging.debug('Encoded: %s', encoded_ps)

    out = remote_cmd(winrm, 'powershell', ['-encodedcommand', encoded_ps])
    if out is None:
        logging.error('Execution of remote script failed')
        return 'Remote script "{}" failed\n'.format(ps)

    return out

####################
# libvirt handling #
####################

def open_libvirt_conn(address):
    try:
        conn = libvirt.open('docker://' + address)
    except Exception:
        logging.error('Connecting libvirt to "%s" failed', address)
        return None

    return conn

def make_libvirt_files(conn, dname, flist):
    domains = conn.listAllDomains()
    with open(dname + 'ContainersInfo.txt', mode = 'wt') as fd:
        out = ''
        for di in domains:
            out += '{} {:5d} {}\n'.format(di.UUIDString(), di.ID(), di.name())
        logging.info('List of domains:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

    for di in domains:
        with open('{}domain-{}.xml'.format(dname, di.UUIDString()),
                  mode = 'wt') as fd:
            out = di.XMLDesc()
            logging.info('XML for domain %s:\n%s', di.UUIDString(), out)
            fd.write(out)
            flist.add(fd.name)

    networks = conn.listAllNetworks()
    for ni in networks:
        with open('{}network-{}.xml'.format(dname, ni.UUIDString()),
                  mode = 'wt') as fd:
            out = ni.XMLDesc()
            logging.info('XML for network %s:\n%s', ni.UUIDString(), out)
            fd.write(out)
            flist.add(fd.name)

    pools = conn.listAllStoragePools(3)
    for pi in pools:
        with open('{}pool-{}.xml'.format(dname, pi.UUIDString()),
                  mode = 'wt') as fd:
            out = pi.XMLDesc()
            logging.info('XML for pool %s:\n%s', pi.UUIDString(), out)
            fd.write(out)
            flist.add(fd.name)

            volumes = pi.listAllVolumes()
            for vi in volumes:
                with open('{}pool-{}-{}.xml'.format(dname, pi.UUIDString(),
                          vi.key().replace('\\', '').replace('?', '')),
                          mode = 'wt') as fd:
                    out = vi.XMLDesc()
                    logging.info('XML for vol %s:\n%s', vi.key(), out)
                    fd.write(out)
                    flist.add(fd.name)

def get_libvirt_info(conn, dname, flist):
    try:
        make_libvirt_files(conn, dname, flist)
    except libvirt.libvirtError as libvirt_e:
        logging.error('Libvirt error: %s', libvirt_e)

    conn.close()

##################
# linux handling #
##################

def get_linux_client(dname, flist):
    cert_dir = get_cert_dir()
    ca_file = os.path.join(cert_dir, 'ca.pem')
    cert_file = os.path.join(cert_dir, 'cert.pem')

    with open(dname + 'Certificates.txt', mode = 'wt') as fd:
        out = cmd_out('openssl x509 -inform PEM -text -fingerprint -noout -in '
                      + ca_file)
        out += cmd_out('openssl x509 -inform PEM -text -fingerprint -noout -in '
                       + cert_file)
        logging.info('Certificates information:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'MoreHostInfo.xml', mode = 'wt') as fd:
        out = cmd_out('uname -a')
        logging.info('More host information:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'AllProcesses.txt', mode = 'wt') as fd:
        out = cmd_out('ps auxf')
        logging.info('List all processes:\n%s', out[:10000] +
                      (out[10000:] and '...\n'))
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'NetConfig.txt', mode = 'wt') as fd:
        out = cmd_out('ifconfig')
        logging.info('Network information:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

####################
# windows handling #
####################

def get_windows_client(dname, flist):
    with open(dname + 'MoreHostInfo.xml', mode = 'wt') as fd:
        out = cmd_out('powershell -Command '
                      '[System.Environment]::OSVersion.VersionString')
        logging.info('More host information: %s', out)
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'AllProcesses.txt', mode = 'wt') as fd:
        out = cmd_out('powershell -Command Get-Process')
        logging.info('List all processes:\n%s', out[:10000] +
                      (out[10000:] and '...\n'))
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'NetConfig.txt', mode = 'wt') as fd:
        out = cmd_out('ipconfig /all')
        logging.info('Network information:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

###########################
# remote windows handling #
###########################

def get_windows_info(wrm, dname, flist):

    with open(dname + 'MoreHostInfo.xml', mode = 'at') as fd:
        out = remote_ps(wrm, '[System.Environment]::OSVersion.VersionString')
        logging.info('More host information: %s', out)
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'AllProcesses.txt', mode = 'at') as fd:
        out = remote_ps(wrm, 'Get-Process')
        logging.info('List all processes:\n%s', out[:10000] +
                      (out[10000:] and '...\n'))
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'NetConfig.txt', mode = 'at') as fd:
        out = remote_cmd(wrm, 'ipconfig', ['/all'])
        logging.info('Network information:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'Certificates.txt', mode = 'at') as fd:
        out = remote_ps(wrm, 'Get-ChildItem "cert:\\LocalMachine\\My"|'
                        'Select-Object *')
        out += remote_ps(wrm, 'Get-ChildItem "cert:\\LocalMachine\\Root"|'
                         'Select-Object *')
        logging.info('Certificates information:\n%s', out[:10000] +
                      (out[10000:] and '...\n'))
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'HostInfo.txt', mode = 'wt') as fd:
        out = remote_ps(wrm, 'Get-ComputerInfo')
        logging.info('Host information:\n%s', out[:10000] +
                      (out[10000:] and '...\n'))
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'InstalledSoftware.txt', mode = 'wt') as fd:
        out = remote_ps(wrm, 'Get-WindowsFeature')
        out += remote_ps(wrm, 'Get-Package | Format-List')
        logging.info('Installed software:\n%s', out[:10000] +
                      (out[10000:] and '...\n'))
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'DockerInfo.txt', mode = 'wt') as fd:
        out = remote_ps(wrm, 'Get-Item (Join-Path $env:ProgramFiles '
                        '"Docker/dockerd.exe")|Format-List Length,VersionInfo')
        out += remote_ps(wrm, 'Get-Item (Join-Path $env:ProgramFiles '
                         '"Docker/joblib.dll")|Format-List Length,VersionInfo')
        out += remote_ps(wrm, 'Get-Content (Join-Path $env:ProgramData '
                         '"docker/config/daemon.json")')
        out += remote_cmd(wrm, 'docker', ['version'])
        out += remote_cmd(wrm, 'docker', ['info'])
        logging.info('Docker information:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'UninstallConfig.txt', mode = 'wt') as fd:
        out = remote_ps(wrm, 'Get-Content (Join-Path $env:ProgramData '
                        '"Virtuozzo/winctfiles/uninst.cfg")')
        logging.info('Uninstall configuration:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

    with open(dname + 'WinLogFiles.txt', mode = 'wt') as fd:
        out = remote_ps(wrm, 'Get-ChildItem -Path (Join-Path $env:ProgramData '
                        '"Virtuozzo/winctfiles/*.log")|Get-Content')
        logging.info('Windows log files:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

    close_winrm(wrm)

########################
# make report handling #
########################

def make_report(dname, flist, options):
    try:
        ver = libvirt.getVersion()
    except Exception:
        ver = 0

    ver = '{}.{}.{}'.format(ver // 1000000, ver // 1000 % 100, ver % 1000)

    with open(dname + 'Report.xml', mode = 'wt') as fd:
        out = report_xml_templet.format(ver, options.description,
                                        options.name, options.email,
                                        options.proxy, options.address)
        logging.info('XML for report:\n%s', out)
        fd.write(out)
        flist.add(fd.name)

    arch = tarfile.open(dname + report_arch, mode = 'w:gz')
    for fi in flist:
        logging.debug('Add file: %s', fi)
        arch.add(fi, arcname = fi.replace(dname, ''))
    arch.close()

########################
# save report handling #
########################

def save_report(dname, flist, options):
    savefile = os.path.abspath(options.reportfile)

    if os.path.isdir(savefile):
        try:
            fd, tmpname = tempfile.mkstemp(dir = savefile,
                                           prefix = 'report_',
                                           suffix = '.tar.gz')
        except Exception:
            logging.error('Can\'t create file in "%s" directory', savefile)
            return False

        logging.debug('File: %s', tmpname)
        os.close(fd)
        savefile = tmpname
    else:
        if os.path.isfile(savefile):
            logging.error('File "%s" already exists', savefile)
            return False

        if not os.path.isdir(os.path.dirname(savefile)):
            logging.error('Directory "%s" doesn\'t exist',
                          os.path.dirname(savefile))
            return False

    try:
        shutil.move(dname + report_arch, savefile)
    except Exception:
        os.remove(savefile)
        logging.error('Moving file to "%s" failed', savefile)
        return False

    logging.critical('Report file: %s', savefile)

    return True

########################
# send report handling #
########################

def send_report(dname, flist, options):
    headers = { 'host' : report_server }
    files = {
        'report' : (report_arch, open(dname + report_arch, 'rb'),
                    'application/octet-stream')
    }

    proxies = None
    if not options.proxy is None:
        proxy = re.split(r'[:@]', options.proxy)
        proxies = { 'https' : 'https://{}:{}@{}:{}/'.format(proxy[0],
            proxy[1], proxy[2], proxy[3]) }

    url = 'https://{}:443/post'.format(report_server)

    urllib3.disable_warnings() # disable SSL warning about host verify
    resp = requests.post(url, files = files, headers = headers,
                         proxies = proxies, verify = False)
    if resp.status_code != requests.codes.ok:
        logging.error('Status code "%s" is %d', url, resp.status_code)
        return False

    report = resp.text.split(';')
    logging.critical('Report ID: %s', report[1])

    return True

###################
# main() handling #
###################

def main():
    options = parse_args()

    log_format = '%(asctime)s %(message)s'
    time_format = '%H:%M:%S'
    logging.basicConfig(format = log_format, datefmt = time_format,
                        level = 30 - options.verbose * 10)

    rc = 0
    flist = set()
    dname = tempfile.mkdtemp() + '/'
    logging.debug('Temp dir: %s', dname)

    if os.name == 'posix':
        get_linux_client(dname, flist)
    elif os.name == 'nt':
        get_windows_client(dname, flist)

    conn = open_libvirt_conn(options.address)
    if conn:
        get_libvirt_info(conn, dname, flist)
    else:
        rc += 1

    wrm = open_winrm(options.address)
    if wrm:
        get_windows_info(wrm, dname, flist)
    else:
        rc += 1

    make_report(dname, flist, options)

    if not options.reportfile is None:
        if not save_report(dname, flist, options):
            logging.error('Saving report failed')
            rc = -1
    else:
        if not send_report(dname, flist, options):
            logging.error('Sending report failed')
            rc = -1

    shutil.rmtree(dname)

    return rc

if __name__ == '__main__':
    try:
        sys.exit(main())
    except SystemExit as sys_e:
        sys.exit(sys_e.code)
