#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# The Qubes OS Project, http://www.qubes-os.org
#
# Copyright (C) 2011  Marek Marczykowski <marmarek@mimuw.edu.pl>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# 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, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#
#
import optparse

import subprocess
import re
import os
import sys
import shutil
import pipes

from optparse import OptionParser
from qubes.qubes import QubesVmCollection, QubesException, system_path
from qubes.qubes import QubesHVm
from qubes.qubes import vm_files
import qubes.imgconverter
from qubes.qubes import vmm

# fields required to be present (and verified) in retrieved desktop file
required_fields = ["Name", "Exec"]

# limits
appmenus_line_size = 1024
appmenus_line_count = 100000

# regexps for sanitization of retrieved values
std_re = re.compile(r"^[/a-zA-Z0-9.,:&()_ -]*$")
fields_regexp = {
    "Name": std_re,
    "GenericName": std_re,
    "Comment": std_re,
    "Categories": re.compile(r"^[a-zA-Z0-9/.;:'() -]*$"),
    "Exec": re.compile(r"^[a-zA-Z0-9()_%&>/{}\"'\\:.= -]*$"),
    "Icon": re.compile(r"^[a-zA-Z0-9/_.-]*$"),
}

CATEGORIES_WHITELIST = {
    # Main Categories
    # http://standards.freedesktop.org/menu-spec/1.1/apa.html 20140507
    'AudioVideo', 'Audio', 'Video', 'Development', 'Education', 'Game',
    'Graphics', 'Network', 'Office', 'Science', 'Settings', 'System',
    'Utility',

    # Additional Categories
    # http://standards.freedesktop.org/menu-spec/1.1/apas02.html
    'Building', 'Debugger', 'IDE', 'GUIDesigner', 'Profiling',
    'RevisionControl', 'Translation', 'Calendar', 'ContactManagement',
    'Database', 'Dictionary', 'Chart', 'Email', 'Finance', 'FlowChart', 'PDA',
    'ProjectManagement', 'Presentation', 'Spreadsheet', 'WordProcessor',
    '2DGraphics', 'VectorGraphics', 'RasterGraphics', '3DGraphics', 'Scanning',
    'OCR', 'Photography', 'Publishing', 'Viewer', 'TextTools',
    'DesktopSettings', 'HardwareSettings', 'Printing', 'PackageManager',
    'Dialup', 'InstantMessaging', 'Chat', 'IRCClient', 'Feed', 'FileTransfer',
    'HamRadio', 'News', 'P2P', 'RemoteAccess', 'Telephony', 'TelephonyTools',
    'VideoConference', 'WebBrowser', 'WebDevelopment', 'Midi', 'Mixer',
    'Sequencer', 'Tuner', 'TV', 'AudioVideoEditing', 'Player', 'Recorder',
    'DiscBurning', 'ActionGame', 'AdventureGame', 'ArcadeGame', 'BoardGame',
    'BlocksGame', 'CardGame', 'KidsGame', 'LogicGame', 'RolePlaying',
    'Shooter', 'Simulation', 'SportsGame', 'StrategyGame', 'Art',
    'Construction', 'Music', 'Languages', 'ArtificialIntelligence',
    'Astronomy', 'Biology', 'Chemistry', 'ComputerScience',
    'DataVisualization', 'Economy', 'Electricity', 'Geography', 'Geology',
    'Geoscience', 'History', 'Humanities', 'ImageProcessing', 'Literature',
    'Maps', 'Math', 'NumericalAnalysis', 'MedicalSoftware', 'Physics',
    'Robotics', 'Spirituality', 'Sports', 'ParallelComputing', 'Amusement',
    'Archiving', 'Compression', 'Electronics', 'Emulator', 'Engineering',
    'FileTools', 'FileManager', 'TerminalEmulator', 'Filesystem', 'Monitor',
    'Security', 'Accessibility', 'Calculator', 'Clock', 'TextEditor',
    'Documentation', 'Adult', 'Core', 'KDE', 'GNOME', 'XFCE', 'GTK', 'Qt',
    'Motif', 'Java', 'ConsoleOnly',

    # Reserved Categories (not whitelisted)
    # http://standards.freedesktop.org/menu-spec/1.1/apas03.html
    # 'Screensaver', 'TrayIcon', 'Applet', 'Shell',
}


def sanitise_categories(untrusted_value):
    untrusted_categories = (c.strip() for c in untrusted_value.split(';') if c)
    categories = (c for c in untrusted_categories if c in CATEGORIES_WHITELIST)

    return ';'.join(categories) + ';'


def fallback_hvm_appmenulist():
    p = subprocess.Popen(["grep", "-rH", "=", "/usr/share/qubes-appmenus/hvm"],
                         stdout=subprocess.PIPE)
    (stdout, stderr) = p.communicate()
    return stdout.splitlines()


def get_appmenus(vm):
    global appmenus_line_count
    global appmenus_line_size
    untrusted_appmenulist = []
    if vm is None:
        while appmenus_line_count > 0:
            untrusted_line = sys.stdin.readline(appmenus_line_size)
            if untrusted_line == "":
                break
            untrusted_appmenulist.append(untrusted_line.strip())
            appmenus_line_count -= 1
        if appmenus_line_count == 0:
            raise QubesException("Line count limit exceeded")
    else:
        p = vm.run('QUBESRPC qubes.GetAppmenus dom0', passio_popen=True,
                   gui=False)
        while appmenus_line_count > 0:
            untrusted_line = p.stdout.readline(appmenus_line_size)
            if untrusted_line == "":
                break
            untrusted_appmenulist.append(untrusted_line.strip())
            appmenus_line_count -= 1
        p.wait()
        if p.returncode != 0:
            if isinstance(vm, QubesHVm):
                untrusted_appmenulist = fallback_hvm_appmenulist()
            else:
                raise QubesException("Error getting application list")
        if appmenus_line_count == 0:
            raise QubesException("Line count limit exceeded")

    appmenus = {}
    line_rx = re.compile(
        r"([a-zA-Z0-9.()_-]+.desktop):([a-zA-Z0-9-]+(?:\[[a-zA-Z@_]+\])?)=(.*)")
    ignore_rx = re.compile(r".*([a-zA-Z0-9._-]+.desktop):(#.*|\s+)$")
    for untrusted_line in untrusted_appmenulist:
        # Ignore blank lines and comments
        if len(untrusted_line) == 0 or ignore_rx.match(untrusted_line):
            continue
        # use search instead of match to skip file path
        untrusted_m = line_rx.search(untrusted_line)
        if untrusted_m:
            filename = untrusted_m.group(1)
            assert '/' not in filename
            assert '\0' not in filename

            untrusted_key = untrusted_m.group(2)
            assert '\0' not in untrusted_key
            assert '\x1b' not in untrusted_key
            assert '=' not in untrusted_key

            untrusted_value = untrusted_m.group(3)
            # TODO add key-dependent asserts

            # Look only at predefined keys
            if untrusted_key in fields_regexp:
                if fields_regexp[untrusted_key].match(untrusted_value):
                    # now values are sanitized
                    key = untrusted_key
                    if key == 'Categories':
                        value = sanitise_categories(untrusted_value)
                    else:
                        value = untrusted_value

                    if filename not in appmenus:
                        appmenus[filename] = {}

                    appmenus[filename][key] = value
                else:
                    print >> sys.stderr, \
                        "Warning: ignoring key %r of %s" % \
                        (untrusted_key, filename)
            # else: ignore this key

    return appmenus


def create_template(path, values):
    # check if all required fields are present
    for key in required_fields:
        if key not in values:
            print >> sys.stderr, "Warning: not creating/updating '%s' " \
                                 "because of missing '%s' key" % (
                                     path, key)
            return

    desktop_entry = ""
    desktop_entry += "[Desktop Entry]\n"
    desktop_entry += "Version=1.0\n"
    desktop_entry += "Type=Application\n"
    desktop_entry += "Terminal=false\n"
    desktop_entry += "X-Qubes-VmName=%VMNAME%\n"

    if 'Icon' in values:
        icon_file = os.path.splitext(os.path.split(path)[1])[0] + '.png'
        desktop_entry += "Icon={0}\n".format(os.path.join(
            '%VMDIR%', vm_files['appmenus_icons_subdir'], icon_file))
    else:
        desktop_entry += "Icon=%XDGICON%\n"

    for key in ["Name", "GenericName"]:
        if key in values:
            desktop_entry += "{0}=%VMNAME%: {1}\n".format(key, values[key])

    # force category X-Qubes-VM
    values["Categories"] = values.get("Categories", "") + "X-Qubes-VM;"

    for key in ["Comment", "Categories"]:
        if key in values:
            desktop_entry += "{0}={1}\n".format(key, values[key])

    desktop_entry += "Exec=qvm-run -q --tray -a %VMNAME% -- {0}\n".format(
        pipes.quote(values['Exec']))
    if not os.path.exists(path) or desktop_entry != open(path, "r").read():
        desktop_file = open(path, "w")
        desktop_file.write(desktop_entry)
        desktop_file.close()


def main():
    env_vmname = os.environ.get("QREXEC_REMOTE_DOMAIN")
    usage = "usage: %prog [options] <vm-name>\n" \
            "Update desktop file templates for given StandaloneVM or TemplateVM"

    parser = OptionParser(usage)
    parser.add_option("-v", "--verbose", action="store_true", dest="verbose",
                      default=False)
    parser.add_option("--force-root", action="store_true", dest="force_root",
                      default=False,
                      help="Force to run, even with root privileges")
    parser.add_option("--force-rpc", action="store_true", dest="force_rpc",
                      default=False,
                      help="Force to start a new RPC call, "
                           "even if called from existing one")
    # Do not use any RPC call, expects data on stdin (in qubes.GetAppmenus
    # format)
    parser.add_option("--offline-mode", dest="offline_mode",
                      action="store_true", default=False,
                      help=optparse.SUPPRESS_HELP)

    (options, args) = parser.parse_args()
    if (len(args) != 1) and env_vmname is None:
        parser.error("You must specify at least the VM name!")

    if env_vmname:
        vmname = env_vmname
    else:
        vmname = args[0]

    if os.geteuid() == 0:
        if not options.force_root:
            print >> sys.stderr, "*** Running this tool as root is strongly " \
                                 "discouraged, this will lead you into " \
                                 "permissions problems."
            print >> sys.stderr, "Retry as unprivileged user."
            print >> sys.stderr, "... or use --force-root to continue anyway."
            exit(1)

    if options.offline_mode:
        vmm.offline_mode = True
    qvm_collection = QubesVmCollection()
    qvm_collection.lock_db_for_reading()
    qvm_collection.load()
    qvm_collection.unlock_db()

    vm = qvm_collection.get_vm_by_name(vmname)

    if vm is None:
        print >> sys.stderr, "ERROR: A VM with the name '{0}' " \
                             "does not exist in the system.".format(
                                 vmname)
        exit(1)

    if vm.template is not None:
        print >> sys.stderr, "ERROR: To sync appmenus for template based VM, " \
                             "do it on template instead"
        exit(1)

    if not options.offline_mode and not vm.is_running():
        print >> sys.stderr, "ERROR: Appmenus can be retrieved only from " \
                             "running VM - start it first"
        exit(1)

    if not options.offline_mode and env_vmname is None or options.force_rpc:
        new_appmenus = get_appmenus(vm)
    else:
        options.verbose = False
        new_appmenus = get_appmenus(None)

    if len(new_appmenus) == 0:
        print >> sys.stderr, "ERROR: No appmenus received, terminating"
        exit(1)

    os.umask(002)

    if not os.path.exists(vm.appmenus_templates_dir):
        os.mkdir(vm.appmenus_templates_dir)

    if not os.path.exists(vm.appmenus_template_icons_dir):
        os.mkdir(vm.appmenus_template_icons_dir)

    # Create new/update existing templates
    if options.verbose:
        print >> sys.stderr, "--> Got {0} appmenus, storing to disk".format(
            str(len(new_appmenus)))
    for appmenu_file in new_appmenus.keys():
        if options.verbose:
            if os.path.exists(
                    os.path.join(vm.appmenus_templates_dir, appmenu_file)):
                print >> sys.stderr, "---> Updating {0}".format(appmenu_file)
            else:
                print >> sys.stderr, "---> Creating {0}".format(appmenu_file)

        # TODO: icons support in offline mode
        if options.offline_mode:
            new_appmenus[appmenu_file].pop('Icon', None)
        if 'Icon' in new_appmenus[appmenu_file]:
            # the following line is used for time comparison
            icondest = os.path.join(vm.appmenus_template_icons_dir,
                                    os.path.splitext(appmenu_file)[0] + '.png')

            try:
                icon = qubes.imgconverter.Image. \
                    get_xdg_icon_from_vm(vm,
                                         new_appmenus[
                                             appmenu_file][
                                             'Icon'])
                if os.path.exists(icondest):
                    old_icon = qubes.imgconverter.Image.load_from_file(icondest)
                else:
                    old_icon = None
                if old_icon is None or icon != old_icon:
                    icon.save(icondest)
            except Exception, e:
                print >> sys.stderr, '----> Failed to get icon for {0}: {1!s}'.\
                    format(appmenu_file, e)

                if os.path.exists(icondest):
                    print >> sys.stderr, '-----> Found old icon, ' \
                                         'using it instead'
                else:
                    del new_appmenus[appmenu_file]['Icon']

        create_template(os.path.join(vm.appmenus_templates_dir, appmenu_file),
                        new_appmenus[appmenu_file])

    # Delete appmenus of removed applications
    if options.verbose:
        print >> sys.stderr, "--> Cleaning old files"
    for appmenu_file in os.listdir(vm.appmenus_templates_dir):
        if not appmenu_file.endswith('.desktop'):
            continue

        if appmenu_file not in new_appmenus:
            if options.verbose:
                print >> sys.stderr, "---> Removing {0}".format(appmenu_file)
            os.unlink(os.path.join(vm.appmenus_templates_dir, appmenu_file))

    if isinstance(vm, QubesHVm):
        if not os.path.exists(os.path.join(vm.appmenus_templates_dir,
                                           os.path.basename(system_path[
                                               'appmenu_start_hvm_template']))):
            shutil.copy(system_path['appmenu_start_hvm_template'],
                        vm.appmenus_templates_dir)

    vm.appmenus_update()
    if hasattr(vm, 'appvms'):
        os.putenv('SKIP_CACHE_REBUILD', '1')
        for child_vm in vm.appvms.values():
            try:
                child_vm.appmenus_update()
            except Exception, e:
                print >> sys.stderr, "---> Failed to recreate appmenus for " \
                                     "'{0}': {1}".format(child_vm.name, str(e))
        subprocess.call(['xdg-desktop-menu', 'forceupdate'])
        if 'KDE_SESSION_UID' in os.environ:
            subprocess.call(['kbuildsycoca' + os.environ.get('KDE_SESSION_VERSION', '4')])
        os.unsetenv('SKIP_CACHE_REBUILD')


main()
