#!/usr/local/munki/munki-python
# encoding: utf-8
#
# Copyright 2010-2025 Greg Neagle.
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
iconimporter

Created by Greg Neagle on 2014-03-03.

Converts and imports icons as png files for Munki repo
"""
from __future__ import absolute_import, print_function

# standard libs
import os
from optparse import OptionParser
import sys

# our libs
from munkilib.cliutils import TempFile
from munkilib.cliutils import pref, path2url


from munkilib import dmgutils
from munkilib import iconutils
from munkilib import munkirepo
from munkilib import osinstaller
from munkilib import osutils
from munkilib import pkgutils

from munkilib.FoundationPlist import (readPlistFromString,
                                      FoundationPlistException)
from munkilib.wrappers import unicode_or_str


def copy_icon_to_repo(repo, name, path):
    '''Copies png file in path to repo as icons/name.png'''
    icon_ref = os.path.join(u'icons', name + u'.png')
    try:
        repo.put_from_local_file(icon_ref, path)
        print(u'\tWrote: %s' % icon_ref)
    except munkirepo.RepoError as err:
        print(u'\tError uploading %s: %s' % (icon_ref, unicode_or_str(err)),
              file=sys.stderr)


def generate_png_from_startosinstall_item(repo, install_item):
    '''Generate a PNG from a disk image containing an Install macOS app'''
    # Since the repo might be a remote repo reached by a web API, we have
    # to download the file first. We might want to extend the Repo plugin
    # "API" to let us get the direct filepath, skipping the need to download for
    # the FileRepo-type repos at least.
    dmg_ref = os.path.join('pkgs', install_item['installer_item_location'])
    dmg_temp = TempFile()
    try:
        repo.get_to_local_file(dmg_ref, dmg_temp.path)
    except munkirepo.RepoError as err:
        print(
            u'\tCan\'t download %s from repo: %s'
            % (dmg_ref, unicode_or_str(err)),
            file=sys.stderr
        )
        return
    mountpoints = dmgutils.mountdmg(dmg_temp.path)
    if mountpoints:
        mountpoint = mountpoints[0]
        app_path = osinstaller.find_install_macos_app(mountpoint)
        if app_path:
            icon_path = iconutils.findIconForApp(app_path)
            if icon_path:
                # use a local temp file to create our png
                icon_temp = TempFile()
                result = iconutils.convertIconToPNG(icon_path, icon_temp.path)
                if result:
                    copy_icon_to_repo(
                        repo, install_item['name'], icon_temp.path)
                else:
                    print(u'\tError converting %s to png.' % icon_path,
                          file=sys.stderr)
            else:
                print(u'\tNo application icons found.')
        else:
            print(u'\tNo application icons found.')
        dmgutils.unmountdmg(mountpoint)


def generate_png_from_dmg_item(repo, install_item):
    '''Generate a PNG from a disk image containing an application'''
    # Since the repo might be a remote repo reached by a web API, we have
    # to download the file first. We might want to extend the Repo plugin
    # "API" to let us get the direct filepath, skipping the need to download for
    # the FileRepo-type repos at least.
    dmg_ref = os.path.join('pkgs', install_item['installer_item_location'])
    dmg_temp = TempFile()
    try:
        repo.get_to_local_file(dmg_ref, dmg_temp.path)
    except munkirepo.RepoError as err:
        print(u'\tCan\'t download %s from repo: %s'
              % (dmg_ref, unicode_or_str(err)), file=sys.stderr)
        return
    mountpoints = dmgutils.mountdmg(dmg_temp.path)
    if mountpoints:
        mountpoint = mountpoints[0]
        apps = [item for item in install_item.get('items_to_copy', [])
                if item.get('source_item', '').endswith('.app')]
        if apps:
            app_path = os.path.join(mountpoint, apps[0]['source_item'])
            icon_path = iconutils.findIconForApp(app_path)
            if icon_path:
                # use a local temp file to create our png
                icon_temp = TempFile()
                result = iconutils.convertIconToPNG(icon_path, icon_temp.path)
                if result:
                    copy_icon_to_repo(
                        repo, install_item['name'], icon_temp.path)
                else:
                    print(u'\tError converting %s to png.' % icon_path,
                          file=sys.stderr)
            else:
                print(u'\tNo application icons found.')
        else:
            print(u'\tNo application icons found.')
        dmgutils.unmountdmg(mountpoint)


def generate_pngs_from_pkg(repo, install_item):
    '''Generate PNGS from applications inside a pkg'''
    icon_paths = []
    mountpoint = None
    pkg_path = None
    item_path = os.path.join(u'pkgs', install_item['installer_item_location'])
    file_temp = TempFile()
    try:
        repo.get_to_local_file(item_path, file_temp.path)
    except munkirepo.RepoError as err:
        print(u'\tCan\'t download %s from repo: %s'
              % (item_path, unicode_or_str(err)), file=sys.stderr)
        return
    if pkgutils.hasValidDiskImageExt(item_path):
        mountpoints = dmgutils.mountdmg(file_temp.path)
        if mountpoints:
            mountpoint = mountpoints[0]
            if install_item.get('package_path'):
                pkg_path = os.path.join(
                    mountpoint, install_item['package_path'])
            else:
                # find first item that appears to be a pkg at the root
                for fileitem in osutils.listdir(mountpoints[0]):
                    if pkgutils.hasValidPackageExt(fileitem):
                        pkg_path = os.path.join(mountpoint, fileitem)
                        break
    elif pkgutils.hasValidPackageExt(item_path):
        pkg_path = file_temp.path

    if pkg_path:
        if os.path.isdir(pkg_path):
            icon_paths = iconutils.extractAppIconsFromBundlePkg(pkg_path)
        else:
            icon_paths = iconutils.extractAppIconsFromFlatPkg(pkg_path)

    if mountpoint:
        dmgutils.unmountdmg(mountpoint)

    if len(icon_paths) == 1:
        icon_temp = TempFile()
        result = iconutils.convertIconToPNG(icon_paths[0], icon_temp.path)
        if result:
            copy_icon_to_repo(
                repo, install_item['name'], icon_temp.path)
    elif len(icon_paths) > 1:
        index = 1
        for icon_path in icon_paths:
            icon_name = install_item['name'] + '_' + str(index)
            icon_temp = TempFile()
            result = iconutils.convertIconToPNG(icon_path, icon_temp.path)
            if result:
                copy_icon_to_repo(repo, icon_name, icon_temp.path)
            index += 1
    else:
        print(u'\tNo application icons found.')


def find_items_to_check(repo, itemlist=None):
    '''Builds a list of items to check; only the latest version
    of an item is retained. If itemlist is given, include items
    only on that list.'''
    try:
        all_catalog_data = repo.get('catalogs/all')
        catalogitems = readPlistFromString(all_catalog_data)
    except (munkirepo.RepoError, FoundationPlistException) as err:
        print('Error getting catalog data from repo: %s' % unicode_or_str(err),
              file=sys.stderr)
        return []

    itemdb = {}
    for catalogitem in catalogitems:
        if itemlist and catalogitem['name'] not in itemlist:
            continue
        name = catalogitem['name']
        if name not in itemdb:
            itemdb[name] = catalogitem
        elif (pkgutils.MunkiLooseVersion(catalogitem['version'])
              > pkgutils.MunkiLooseVersion(itemdb[name]['version'])):
            itemdb[name] = catalogitem
    pkg_list = []
    for key in itemdb:
        pkg_list.append(itemdb[key])
    return pkg_list


def generate_pngs_from_munki_items(repo, force=False, itemlist=None):
    '''Generate PNGs from either pkgs or disk images containing applications'''
    itemlist = find_items_to_check(repo, itemlist=itemlist)
    try:
        icons_list = repo.itemlist('icons')
    except munkirepo.RepoError:
        icons_list = []
    for item in itemlist:
        print(u'Processing %s...' % item['name'])
        icon_name = item.get('icon_name') or item['name']
        if not os.path.splitext(icon_name)[1]:
            icon_name += u'.png'
        if icon_name in icons_list and not force:
            print(u'Found existing icon at %s' % icon_name)
            continue
        installer_type = item.get('installer_type')
        if installer_type == 'copy_from_dmg':
            generate_png_from_dmg_item(repo, item)
        elif installer_type == 'startosinstall':
            generate_png_from_startosinstall_item(repo, item)
        elif installer_type in [None, '']:
            generate_pngs_from_pkg(repo, item)
        else:
            print(u'\tCan\'t process installer_type: %s' % installer_type)


def main():
    '''Main'''
    usage = "usage: %prog [options] [/path/to/repo_root]"
    parser = OptionParser(usage=usage)
    parser.add_option(
        '--force', '-f', action='store_true', dest='force',
        help='Create pngs even if there is an existing icon in the repo.')
    parser.add_option(
        '--item', '-i', action='append', type='string', dest='items',
        help='Only run for given pkginfo item name(s).')
    parser.add_option('--plugin', default=pref('plugin'),
                      help='Optional. Custom plugin to connect to repo.')
    parser.add_option('--repo_url', '--repo-url', default=pref('repo_url'),
                      help='Optional repo fileshare URL used by repo plugin.')
    parser.set_defaults(force=False)
    options, arguments = parser.parse_args()

    # Make sure we have a path to work with
    if arguments:
        repo_path = arguments[0].rstrip("/")
    else:
        repo_path = pref('repo_path')
    if not options.repo_url and repo_path:
        options.repo_url = path2url(repo_path)

    if options.plugin is None:
        options.plugin = 'FileRepo'

    # Make sure the repo exists
    try:
        repo = munkirepo.connect(options.repo_url, options.plugin)
    except munkirepo.RepoError as err:
        print(u'Could not connect to munki repo: %s' % unicode_or_str(err),
              file=sys.stderr)
        exit(-1)

    # generate icons!
    generate_pngs_from_munki_items(
        repo, force=options.force, itemlist=options.items)

    # clean up
    osutils.cleanUpTmpDir()

if __name__ == '__main__':
    main()