#! /usr/bin/python3
# ----------------------------------------------------------------------
#    Copyright (C) 2013 Kshitij Gupta <kgupta8592@gmail.com>
#    Copyright (C) 2014-2016 Christian Boltz <apparmor@cboltz.de>
#
#    This program is free software; you can redistribute it and/or
#    modify it under the terms of version 2 of the GNU General Public
#    License as published by the Free Software Foundation.
#
#    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.
#
# ----------------------------------------------------------------------
import argparse
import os

import apparmor.aa
import apparmor.aamode

import apparmor.severity
import apparmor.cleanprofile as cleanprofile
import apparmor.ui as aaui

from apparmor.aa import (add_to_options, available_buttons, combine_name, delete_duplicates,
                         get_profile_filename, is_known_rule, match_includes, profile_storage,
                         set_options_audit_mode, propose_file_rules, selection_to_rule_obj)
from apparmor.aare import AARE
from apparmor.common import AppArmorException
from apparmor.regex import re_match_include


# setup exception handling
from apparmor.fail import enable_aa_exception_handler
enable_aa_exception_handler()

# setup module translations
from apparmor.translations import init_translation
_ = init_translation()

parser = argparse.ArgumentParser(description=_('Merge the given profiles into /etc/apparmor.d/ (or the directory specified with -d)'))
parser.add_argument('files', nargs='+', type=str, help=_('Profile(s) to merge'))
#parser.add_argument('other', nargs='?', type=str, help=_('other profile'))
parser.add_argument('-d', '--dir', type=str, help=_('path to profiles'))
#parser.add_argument('-a', '--auto', action='store_true', help=_('Automatically merge profiles, exits incase of *x conflicts'))
args = parser.parse_args()

args.other = None
# 2-way merge or 3-way merge based on number of params
merge_mode = 2 #if args.other == None else 3

profiles = [args.files, [args.other]]

profiledir = args.dir
if profiledir:
    apparmor.aa.profile_dir = apparmor.aa.get_full_path(profiledir)
    if not os.path.isdir(apparmor.aa.profile_dir):
        raise AppArmorException(_("%s is not a directory.") %profiledir)

def reset_aa():
    apparmor.aa.aa = apparmor.aa.hasher()
    apparmor.aa.filelist = apparmor.aa.hasher()
    apparmor.aa.include = dict()
    apparmor.aa.existing_profiles = apparmor.aa.hasher()
    apparmor.aa.original_aa = apparmor.aa.hasher()

def find_profiles_from_files(files):
    profile_to_filename = dict()
    for file_name in files:
        apparmor.aa.read_profile(file_name, True)
        for profile_name in apparmor.aa.filelist[file_name]['profiles'].keys():
            profile_to_filename[profile_name] = file_name
        reset_aa()

    return profile_to_filename

def find_files_from_profiles(profiles):
    profile_to_filename = dict()
    apparmor.aa.read_profiles()

    for profile_name in profiles:
        profile_to_filename[profile_name] = apparmor.aa.get_profile_filename(profile_name)

    reset_aa()

    return profile_to_filename

def main():
    profiles_to_merge = set()

    base_files, other_files = profiles

    base_profile_to_file = find_profiles_from_files(base_files)

    profiles_to_merge = profiles_to_merge.union(set(base_profile_to_file.keys()))

    other_profile_to_file = dict()

    if merge_mode == 3:
        other_profile_to_file = find_profiles_from_files(other_files)
        profiles_to_merge.add(other_profile_to_file.keys())

    user_profile_to_file = find_files_from_profiles(profiles_to_merge)

#    print(base_files,"\n",other_files)
#    print(base_profile_to_file,"\n",other_profile_to_file,"\n",user_profile_to_file)
#    print(profiles_to_merge)

    for profile_name in profiles_to_merge:
        aaui.UI_Info("\n\n" + _("Merging profile for %s" % profile_name))
        user_file = user_profile_to_file[profile_name]
        base_file = base_profile_to_file.get(profile_name, None)
        other_file =  None

        if merge_mode == 3:
            other_file = other_profile_to_file.get(profile_name, None)

        if base_file == None:
            if other_file == None:
                continue

            act([user_file, other_file, None], 2, profile_name)
        else:
            if other_file == None:
                act([user_file, base_file, None], 2, profile_name)
            else:
                act([user_file, base_file, other_file], 3, profile_name)

        reset_aa()

def act(files, merge_mode, merging_profile):
    mergeprofiles = Merge(files)
    #Get rid of common/superfluous stuff
    mergeprofiles.clear_common()

#    if not args.auto:
    if 1 == 1:  # workaround to avoid lots of whitespace changes
        if merge_mode == 3:
            mergeprofiles.ask_the_questions('other', merging_profile)

            mergeprofiles.clear_common()

        mergeprofiles.ask_the_questions('base', merging_profile)

        q = aaui.PromptQuestion()
        q.title = _('Changed Local Profiles')
        q.explanation = _('The following local profiles were changed. Would you like to save them?')
        q.functions = ['CMD_SAVE_CHANGES', 'CMD_VIEW_CHANGES', 'CMD_ABORT', 'CMD_IGNORE_ENTRY']
        q.default = 'CMD_VIEW_CHANGES'
        q.options = [merging_profile]
        q.selected = 0

        ans = ''
        arg = None
        programs = list(mergeprofiles.user.aa.keys())
        program = programs[0]
        while ans != 'CMD_SAVE_CHANGES':
            ans, arg = q.promptUser()
            if ans == 'CMD_SAVE_CHANGES':
                apparmor.aa.write_profile_ui_feedback(program)
                apparmor.aa.reload_base(program)
            elif ans == 'CMD_VIEW_CHANGES':
                for program in programs:
                    apparmor.aa.original_aa[program] = apparmor.aa.deepcopy(apparmor.aa.aa[program])
                #oldprofile = apparmor.serialize_profile(apparmor.original_aa[program], program, '')
                newprofile = apparmor.aa.serialize_profile(mergeprofiles.user.aa[program], program, '')
                apparmor.aa.display_changes_with_comments(mergeprofiles.user.filename, newprofile)
            elif ans == 'CMD_IGNORE_ENTRY':
                break


class Merge(object):
    def __init__(self, profiles):
        user, base, other = profiles

        #Read and parse base profile and save profile data, include data from it and reset them
        apparmor.aa.read_profile(base, True)
        self.base = cleanprofile.Prof(base)

        reset_aa()

        #Read and parse other profile and save profile data, include data from it and reset them
        if merge_mode == 3:
            apparmor.aa.read_profile(other, True)
            self.other = cleanprofile.Prof(other)
            reset_aa()

        #Read and parse user profile
        apparmor.aa.read_profile(user, True)
        self.user = cleanprofile.Prof(user)

    def clear_common(self):
        deleted = 0

        if merge_mode == 3:
            #Remove off the parts in other profile which are common/superfluous from user profile
            user_other = cleanprofile.CleanProf(False, self.user, self.other)
            deleted += user_other.compare_profiles()

        #Remove off the parts in base profile which are common/superfluous from user profile
        user_base = cleanprofile.CleanProf(False, self.user, self.base)
        deleted += user_base.compare_profiles()

        if merge_mode == 3:
            #Remove off the parts in other profile which are common/superfluous from base profile
            base_other = cleanprofile.CleanProf(False, self.base, self.other)
            deleted += base_other.compare_profiles()

    def ask_conflict_mode(self, profile, hat, old_profile, merge_profile):
        '''ask user about conflicting exec rules'''
        for oldrule in old_profile['file'].rules:
            conflictingrules = merge_profile['file'].get_exec_conflict_rules(oldrule)

            if conflictingrules.rules:
                q = aaui.PromptQuestion()
                q.headers = [_('Path'), oldrule.path.regex]
                q.headers += [_('Select the appropriate mode'), '']
                options = []
                options.append(oldrule.get_clean())
                for rule in conflictingrules.rules:
                    options.append(rule.get_clean())
                q.options = options
                q.functions = ['CMD_ALLOW', 'CMD_ABORT']
                done = False
                while not done:
                    ans, selected = q.promptUser()
                    if ans == 'CMD_ALLOW':
                        if selected == 0:
                            pass  # just keep the existing rule
                        elif selected > 0:
                            # replace existing rule with merged one
                            old_profile['file'].delete(oldrule)
                            old_profile['file'].add(conflictingrules.rules[selected - 1])
                        else:
                            raise AppArmorException(_('Unknown selection'))

                        for rule in conflictingrules.rules:
                            merge_profile['file'].delete(rule)  # make sure aa-mergeprof doesn't ask to add conflicting rules later

                        done = True

    def ask_the_questions(self, other, profile):
        aa = self.user.aa  # keep references so that the code in this function can use the short name
        changed = apparmor.aa.changed  # (and be more in sync with aa.py ask_the_questions())

        if other == 'other':
            other = self.other
        else:
            other = self.base
        #print(other.aa)

        #Add the file-wide includes from the other profile to the user profile
        apparmor.aa.loadincludes()
        done = False

        options = []
        for inc in other.filelist[other.filename]['include'].keys():
            if not inc in self.user.filelist[self.user.filename]['include'].keys():
                options.append('#include <%s>' %inc)

        default_option = 1

        q = aaui.PromptQuestion()
        q.options = options
        q.selected = default_option - 1
        q.headers = [_('File includes'), _('Select the ones you wish to add')]
        q.functions = ['CMD_ALLOW', 'CMD_IGNORE_ENTRY', 'CMD_ABORT', 'CMD_FINISHED']
        q.default = 'CMD_ALLOW'

        while not done and options:
            ans, selected = q.promptUser()
            if ans == 'CMD_IGNORE_ENTRY':
                done = True
            elif ans == 'CMD_ALLOW':
                selection = options[selected]
                inc = re_match_include(selection)
                self.user.filelist[self.user.filename]['include'][inc] = True
                options.pop(selected)
                aaui.UI_Info(_('Adding %s to the file.') % selection)
            elif ans == 'CMD_FINISHED':
                return

        sev_db = apparmor.aa.sev_db
        if not sev_db:
            sev_db = apparmor.severity.Severity(apparmor.aa.CONFDIR + '/severity.db', _('unknown'))

        sev_db.unload_variables()
        sev_db.load_variables(get_profile_filename(profile))

        for hat in sorted(other.aa[profile].keys()):

            if not aa[profile].get(hat):
                ans = ''
                while ans not in ['CMD_ADDHAT', 'CMD_ADDSUBPROFILE', 'CMD_DENY']:
                    q = aaui.PromptQuestion()
                    q.headers += [_('Profile'), profile]

                    if other.aa[profile][hat]['profile']:
                        q.headers += [_('Requested Subprofile'), hat]
                        q.functions.append('CMD_ADDSUBPROFILE')
                    else:
                        q.headers += [_('Requested Hat'), hat]
                        q.functions.append('CMD_ADDHAT')

                    q.functions += ['CMD_DENY', 'CMD_ABORT', 'CMD_FINISHED']

                    q.default = 'CMD_DENY'

                    ans = q.promptUser()[0]

                    if ans == 'CMD_FINISHED':
                        return

                if ans == 'CMD_DENY':
                    continue  # don't ask about individual rules if the user doesn't want the additional subprofile/hat

                if other.aa[profile][hat]['profile']:
                    aa[profile][hat] = profile_storage(profile, hat, 'mergeprof ask_the_questions() - missing subprofile')
                    aa[profile][hat]['profile'] = True
                else:
                    aa[profile][hat] = profile_storage(profile, hat, 'mergeprof ask_the_questions() - missing hat')
                    aa[profile][hat]['profile'] = False

            #Add the includes from the other profile to the user profile
            done = False

            options = []
            for inc in other.aa[profile][hat]['include'].keys():
                if not inc in aa[profile][hat]['include'].keys():
                    options.append('#include <%s>' %inc)

            default_option = 1

            q = aaui.PromptQuestion()
            q.options = options
            q.selected = default_option - 1
            q.headers = [_('File includes'), _('Select the ones you wish to add')]
            q.functions = ['CMD_ALLOW', 'CMD_IGNORE_ENTRY', 'CMD_ABORT', 'CMD_FINISHED']
            q.default = 'CMD_ALLOW'

            while not done and options:
                ans, selected = q.promptUser()
                if ans == 'CMD_IGNORE_ENTRY':
                    done = True
                elif ans == 'CMD_ALLOW':
                    selection = options[selected]
                    inc = re_match_include(selection)
                    deleted = apparmor.aa.delete_duplicates(aa[profile][hat], inc)
                    aa[profile][hat]['include'][inc] = True
                    options.pop(selected)
                    aaui.UI_Info(_('Adding %s to the file.') % selection)
                    if deleted:
                        aaui.UI_Info(_('Deleted %s previous matching profile entries.') % deleted)
                elif ans == 'CMD_FINISHED':
                    return

            # check for and ask about conflicting exec modes
            self.ask_conflict_mode(profile, hat, aa[profile][hat], other.aa[profile][hat])

            for ruletype in apparmor.aa.ruletypes:
                if other.aa[profile][hat].get(ruletype, False): # needed until we have proper profile initialization
                    for rule_obj in other.aa[profile][hat][ruletype].rules:

                        if is_known_rule(aa[profile][hat], ruletype, rule_obj):
                            continue

                        default_option = 1
                        options = []
                        newincludes = match_includes(aa[profile][hat], ruletype, rule_obj)
                        q = aaui.PromptQuestion()
                        if newincludes:
                            options += list(map(lambda inc: '#include <%s>' % inc, sorted(set(newincludes))))

                        if ruletype == 'file' and rule_obj.path:
                            options += propose_file_rules(aa[profile][hat], rule_obj)
                        else:
                            options.append(rule_obj.get_clean())

                        done = False
                        while not done:
                            q.options = options
                            q.selected = default_option - 1
                            q.headers = [_('Profile'), combine_name(profile, hat)]
                            q.headers += rule_obj.logprof_header()

                            # Load variables into sev_db? Not needed/used for capabilities and network rules.
                            severity = rule_obj.severity(sev_db)
                            if severity != sev_db.NOT_IMPLEMENTED:
                                q.headers += [_('Severity'), severity]

                            q.functions = available_buttons(rule_obj)
                            q.default = q.functions[0]

                            ans, selected = q.promptUser()
                            selection = options[selected]
                            if ans == 'CMD_IGNORE_ENTRY':
                                done = True
                                break

                            elif ans == 'CMD_FINISHED':
                                return

                            elif ans.startswith('CMD_AUDIT'):
                                if ans == 'CMD_AUDIT_NEW':
                                    rule_obj.audit = True
                                    rule_obj.raw_rule = None
                                else:
                                    rule_obj.audit = False
                                    rule_obj.raw_rule = None

                                options = set_options_audit_mode(rule_obj, options)

                            elif ans == 'CMD_ALLOW':
                                done = True
                                changed[profile] = True

                                inc = re_match_include(selection)
                                if inc:
                                    deleted = delete_duplicates(aa[profile][hat], inc)

                                    aa[profile][hat]['include'][inc] = True

                                    aaui.UI_Info(_('Adding %s to profile.') % selection)
                                    if deleted:
                                        aaui.UI_Info(_('Deleted %s previous matching profile entries.') % deleted)

                                else:
                                    rule_obj = selection_to_rule_obj(rule_obj, selection)
                                    deleted = aa[profile][hat][ruletype].add(rule_obj, cleanup=True)

                                    aaui.UI_Info(_('Adding %s to profile.') % rule_obj.get_clean())
                                    if deleted:
                                        aaui.UI_Info(_('Deleted %s previous matching profile entries.') % deleted)

                            elif ans == 'CMD_DENY':
                                if re_match_include(selection):
                                    aaui.UI_Important("Denying via an include file isn't supported by the AppArmor tools")

                                else:
                                    done = True
                                    changed[profile] = True

                                    rule_obj = selection_to_rule_obj(rule_obj, selection)
                                    rule_obj.deny = True
                                    rule_obj.raw_rule = None  # reset raw rule after manually modifying rule_obj
                                    deleted = aa[profile][hat][ruletype].add(rule_obj, cleanup=True)
                                    aaui.UI_Info(_('Adding %s to profile.') % rule_obj.get_clean())
                                    if deleted:
                                        aaui.UI_Info(_('Deleted %s previous matching profile entries.') % deleted)

                            elif ans == 'CMD_GLOB':
                                if not re_match_include(selection):
                                    globbed_rule_obj = selection_to_rule_obj(rule_obj, selection)
                                    globbed_rule_obj.glob()
                                    options, default_option = add_to_options(options, globbed_rule_obj.get_raw())

                            elif ans == 'CMD_GLOBEXT':
                                if not re_match_include(selection):
                                    globbed_rule_obj = selection_to_rule_obj(rule_obj, selection)
                                    globbed_rule_obj.glob_ext()
                                    options, default_option = add_to_options(options, globbed_rule_obj.get_raw())

                            elif ans == 'CMD_NEW':
                                if not re_match_include(selection):
                                    edit_rule_obj = selection_to_rule_obj(rule_obj, selection)
                                    prompt, oldpath = edit_rule_obj.edit_header()

                                    newpath = aaui.UI_GetString(prompt, oldpath)
                                    if newpath:
                                        try:
                                            input_matches_path = rule_obj.validate_edit(newpath)  # note that we check against the original rule_obj here, not edit_rule_obj (which might be based on a globbed path)
                                        except AppArmorException:
                                            aaui.UI_Important(_('The path you entered is invalid (not starting with / or a variable)!'))
                                            continue

                                        if not input_matches_path:
                                            ynprompt = _('The specified path does not match this log entry:\n\n  Log Entry: %(path)s\n  Entered Path:  %(ans)s\nDo you really want to use this path?') % { 'path': oldpath, 'ans': newpath }
                                            key = aaui.UI_YesNo(ynprompt, 'n')
                                            if key == 'n':
                                                continue

                                        edit_rule_obj.store_edit(newpath)
                                        options, default_option = add_to_options(options, edit_rule_obj.get_raw())
                                        apparmor.aa.user_globs[newpath] = AARE(newpath, True)

                            else:
                                done = False

if __name__ == '__main__':
    main()
