#! /usr/bin/python3
# ----------------------------------------------------------------------
#    Copyright (C) 2025 Maxime Bélair <maxime.belair@canonical.com>
#
#    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 glob
import re
import json
import sys
import os

from apparmor import aa
from apparmor.common import AppArmorException
from apparmor.translations import init_translation
from apparmor.regex import expand_path_braces, resolve_variables
from apparmor.aare import convert_regexp

_ = init_translation()

MAX_RECURSION = 10


def glob_superset(pattern):
    """Build a glob whose matches are a superset: replace each remaining
    (within-segment) brace group with '*', keeping '*', '?', '[...]' and '**'.
    The exact regex discards the over-matches afterwards."""
    out, i = [], 0
    while i < len(pattern):
        if pattern[i] == '{':
            depth = 0
            for j in range(i, len(pattern)):
                if pattern[j] == '{':
                    depth += 1
                elif pattern[j] == '}':
                    depth -= 1
                    if depth == 0:
                        break
            else:
                out.append(pattern[i:])
                break
            out.append('*')
            i = j + 1
        else:
            out.append(pattern[i])
            i += 1
    return ''.join(out)


def has_matching_file(pattern, xattrs=None):
    """Return the first existing non-symlinked path the attachment matches, or
    None. The kernel matches attachments as a regex, so glob a cheap superset of
    candidates and keep those matching the exact regex without enumerating the
    {X,} character runs (which could expand to billions of strings)."""
    regex = re.compile(convert_regexp(pattern))
    try:
        variants = expand_path_braces(pattern)
    except AppArmorException as e:
        # Only a pathological pattern reaches this, should never happen.
        sys.stderr.write(_('aa-show-usage: {}\n').format(e.value))
        sys.exit(-1)
    for variant in variants:
        for path in glob.iglob(glob_superset(variant), recursive=True):
            if not regex.match(path):
                continue
            if os.path.realpath(path) != os.path.abspath(path):  # remove symlinks
                continue
            if not xattrs or all(os.getxattr(path, name) == val for name, val in xattrs.items()):
                return path
    return None


def display_profile_text(used, unused, show_matching_path):
    if used:
        print(_('Used profiles:'))
        for (name, attach, path, match) in used:
            print(_('  Profile {} for {} ({}) {}').format(name, attach, path, ('→ ' + match) if show_matching_path else ''))
    if unused:
        print(_('Unused profiles:'))
        for (name, attach, path, match) in unused:
            print(_('  Profile {} for {} ({}) ').format(name, attach, path))


def profiles_to_json(profiles):
    result = []
    for profile_name, attach, path, matching_path in profiles:
        entry = {'name': profile_name, 'attach': attach, 'path': path}
        if matching_path:
            entry['matching_path'] = matching_path
        result.append(entry)
    return result


def display_profile_json(used, unused):
    profiles = {}
    profiles['version'] = 1  # JSON format version - increase if you change the json structure
    if used:
        profiles['used'] = profiles_to_json(used)
    if unused:
        profiles['unused'] = profiles_to_json(unused)

    print(json.dumps(profiles, indent=2))


def filter_profile(path, profile_name, attach, prof_filter):
    if prof_filter['flags'] and not prof_filter['flags'].match(aa.active_profiles.profiles[profile_name].data['flags'] or ''):
        return False
    if prof_filter['name'] and not prof_filter['name'].match(profile_name or ''):
        return False
    if prof_filter['attach'] and not prof_filter['attach'].match(attach or ''):
        return False
    if prof_filter['path'] and not prof_filter['path'].match(path or ''):
        return False

    return True


def get_used_profiles(args, prof_filter):
    aa.init_aa(confdir=args.configdir or os.getenv('__AA_CONFDIR'), profiledir=args.dir)
    aa.read_profiles()
    used = []
    unused = []

    for a, v in aa.active_profiles.attachments.items():
        filename = v['f']
        profile_name = v['p']
        if not filter_profile(filename, profile_name, a, prof_filter):
            continue

        var_dict = aa.active_profiles.get_all_merged_variables(filename, aa.include_list_recursive(aa.active_profiles.files[filename], True))
        resolved = resolve_variables(a, var_dict)
        matching_path = None
        for entry in resolved:
            matching_path = has_matching_file(entry)
            if matching_path:
                break

        if matching_path and args.show_type != 'unused':
            used.append((profile_name, a, filename, matching_path))
        if not matching_path and args.show_type != 'used':
            unused.append((profile_name, a, filename, matching_path))

    return used, unused


def main():
    parser = argparse.ArgumentParser(description=_('Check which profiles are used'))
    parser.add_argument('-s', '--show-type', type=str, default='all', choices=['all', 'used', 'unused'], help=_('Type of profiles to show'))
    parser.add_argument('-j', '--json', action='store_true', help=_('Output in JSON'))
    parser.add_argument('-d', '--dir', type=str, help=_('Path to profiles'))
    parser.add_argument('--show-matching-path', action='store_true', help=_('Show the path of a file matching the profile'))
    parser.add_argument('--configdir', type=str, help=argparse.SUPPRESS)

    filter_group = parser.add_argument_group(_('Filtering options'),
                                             description=(_('Filters are used to reduce the output of information to only '
                                                            'those entries that will match the filter. Filters use Python\'s regular '
                                                            'expression syntax.')))
    filter_group.add_argument('--filter.flags', dest='filter_flags', metavar='FLAGS', help=_('Filter by flags'))
    filter_group.add_argument('--filter.profile_name', dest='filter_name', metavar='PROFILE_NAME', help=_('Filter by profile name'))
    filter_group.add_argument('--filter.profile_attach', dest='filter_attach', metavar='PROFILE_ATTACH', help=_('Filter by profile attachment'))
    filter_group.add_argument('--filter.profile_path', dest='filter_path', metavar='PROFILE_PATH', help=_('Filter by profile path'))

    # If not a TTY then assume running in test mode and fix output width
    if not sys.stdout.isatty():
        parser.formatter_class = lambda prog: argparse.HelpFormatter(prog, width=80)

    args = parser.parse_args()

    prof_filter = {
        'flags':    re.compile(args.filter_flags) if args.filter_flags else None,
        'name':     re.compile(args.filter_name) if args.filter_name else None,
        'attach':   re.compile(args.filter_attach) if args.filter_attach else None,
        'path':     re.compile(args.filter_path) if args.filter_path else None,
    }

    used, unused = get_used_profiles(args, prof_filter)

    if args.json:
        display_profile_json(used, unused)
    else:
        display_profile_text(used, unused, args.show_matching_path)


if __name__ == '__main__':
    main()
