#!/usr/bin/env python
from config import TRACKER_PASSWORD_LENGTH_MIN, TRACKER_PASSWORD_LENGTH_MAX
from app.pacman import update as update_pacman_db, search, get_configpath, archs, primary_arch
from app import db
from app.model import CVEGroup, CVEGroupPackage, CVE, CVEGroupEntry, Package
from app.model.user import User, username_regex
from app.model.enum import Status, UserRole, affected_to_status, status_to_affected, highest_severity
from app.user import random_string, hash_password
from sqlalchemy import func
from collections import defaultdict
from argparse import ArgumentParser
from sys import argv, exit
from subprocess import check_output
from getpass import getpass
from re import match, IGNORECASE
from datetime import datetime


def update_group_status():
    groups = (db.session.query(CVEGroup, func.group_concat(CVEGroupPackage.pkgname, ' '))
              .join(CVEGroupPackage)
              .filter(CVEGroup.status.in_([Status.vulnerable, Status.testing]))
              .group_by(CVEGroupPackage.group_id)).all()
    for group, pkgnames in groups:
        pkgnames = pkgnames.split(' ')
        new_status = affected_to_status(status_to_affected(group.status), pkgnames[0], group.fixed)
        if group.status is not new_status:
            print('    --> Updated {}: {} -> {}'.format(group.name, group.status, new_status))
        group.status = new_status
    db.session.commit()


def recalc_group_status():
    groups = (db.session.query(CVEGroup, func.group_concat(CVEGroupPackage.pkgname, ' '))
              .join(CVEGroupPackage)
              .group_by(CVEGroupPackage.group_id)).all()
    for group, pkgnames in groups:
        pkgnames = pkgnames.split(' ')
        new_status = affected_to_status(status_to_affected(group.status), pkgnames[0], group.fixed)
        if group.status is not new_status:
            print('    --> Updated {}: {} -> {}'.format(group.name, group.status, new_status))
        group.status = new_status
    db.session.commit()


def recalc_group_severity():
    entries = (db.session.query(CVEGroup, CVEGroupEntry, CVE)
               .join(CVEGroupEntry).join(CVE)
               .group_by(CVEGroupEntry.group_id).group_by(CVE.id)).all()
    issues = defaultdict(set)
    for group, entry, issue in entries:
        issues[group].add(issue)
    for group, issues in issues.items():
        new_severity = highest_severity([issue.severity for issue in issues])
        if group.severity is not new_severity:
            print('    --> Updated {}: {} -> {}'.format(group.name, group.severity, new_severity))
        group.severity = new_severity
    db.session.commit()


def update_package_cache():
    print('    --> Querying alpm database...')
    packages = search('', filter_duplicate_packages=False, sort_results=False)
    pkgbases = {}
    packages_by_arch = defaultdict(set)
    # assign packages per arch list (unify any into 64bit)
    for package in packages:
        arch = package.arch if package.arch != 'any' else primary_arch
        packages_by_arch[arch].add(package.name)

    print('    --> Collecting pkgbases...')
    # query all pkgbases per arch
    for arch in archs:
        cmd = ['expac', '-t', '%s', '-S', '--config', get_configpath(arch), '%n %e %v %b']
        cmd.extend(list(packages_by_arch[arch]))
        bases = check_output(cmd).decode().split('\n')[:-1]
        latest = None
        for base in bases:
            pkgname, pkgbase, version, builddate = (base.split(' '))
            pkgbaseorname = pkgbase if pkgbase != '(null)' else pkgname
            pkgbases[pkgname] = pkgbaseorname
            if not latest or builddate > latest['builddate']:
                latest = dict(pkg=pkgbaseorname, version=version, builddate=builddate)
    latest_package = 'None' if not latest else \
                     '{} {} {}'.format(latest['pkg'], latest['version'],
                                       datetime.fromtimestamp(int(latest['builddate'])).strftime('%c'))
    print('    --> Latest package: {}'.format(latest_package))

    print('    --> Updating database cache...')
    new_packages = []
    for package in packages:
        new_packages.append({
            'name': package.name,
            'base': pkgbases[package.name],
            'version': package.version,
            'description': package.desc,
            'url': package.url,
            'arch': package.arch,
            'database': package.db.name,
            'filename': package.filename,
            'md5sum': package.md5sum,
            'sha256sum': package.sha256sum,
            'builddate': package.builddate
        })
    Package.query.delete()
    db.session.bulk_insert_mappings(Package, new_packages)
    db.session.commit()


def db_vacuum():
    db.session.execute('VACUUM')


def create_user():
    user = User()
    user.active = True

    print('Username: ', end='')
    user.name = input()
    if not user.name or not match(username_regex, user.name) or len(user.name) > User.NAME_LENGTH:
        print('ERROR: Invalid username')
        exit(1)

    user.password = getpass()
    if not user.password:
        user.password = random_string()
        print('Generated password: {}'.format(user.password))
    if len(user.password) > TRACKER_PASSWORD_LENGTH_MAX or len(user.password) < TRACKER_PASSWORD_LENGTH_MIN:
        print('ERROR: Password must be between {} and {} characters.'
              .format(TRACKER_PASSWORD_LENGTH_MIN, TRACKER_PASSWORD_LENGTH_MAX))
        exit(1)
    user.salt = random_string()
    user.password = hash_password(user.password, user.salt)

    print('E-Mail: ', end='')
    user.email = input()
    if not match(r'^.+@([^.@][^@]+)$', user.email, IGNORECASE):
        print('ERROR: Invalid E-Mail')
        exit(1)

    print('Role: ', end='')
    user.role = UserRole.fromstring(input())
    if not user.role:
        user.role = UserRole.reporter

    db.session.add(user)
    db.session.commit()


if __name__ == "__main__":
    no_args = 1 >= len(argv)
    parser = ArgumentParser(prog='update')
    parser.add_argument('--create-user', action='store_true')
    parser.add_argument('--pacman-db', action='store_true', default=no_args)
    parser.add_argument('--group-status', action='store_true', default=no_args)
    parser.add_argument('--recalc-group-status', action='store_true')
    parser.add_argument('--recalc-group-severity', action='store_true')
    parser.add_argument('--package-cache', action='store_true', default=no_args)
    parser.add_argument('--db-vacuum', action='store_true')
    args = parser.parse_args()

    if args.create_user:
        print("[+] Creating user...")
        create_user()
        exit(0)

    if args.pacman_db:
        print("[+] Update pacman db...")
        update_pacman_db(force=True)

    if args.package_cache:
        print("[+] Update package cache...")
        update_package_cache()

    if args.group_status:
        print("[+] Update group status...")
        update_group_status()

    if args.recalc_group_status:
        print("[+] Recalc group status...")
        recalc_group_status()

    if args.recalc_group_severity:
        print("[+] Recalc group severity...")
        recalc_group_severity()

    if args.db_vacuum:
        print("[+] Compact database...")
        db_vacuum()
