#!/usr/bin/env python
#
#    migrate-pax: this file is part of the elfix package
#    Copyright (C) 2012  Anthony G. Basile
#
#    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 3 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, see <http://www.gnu.org/licenses/>.


# We use portage's NEEDED.ELF.2 file.  The format is in
# /usr/lib/portage/bin/misc-functions.sh ~line 520
# echo "${arch:3};${obj};${soname};${rpath};${needed}" \
# >> "${PORTAGE_BUILDDIR}"/build-info/NEEDED.ELF.2

import os
import re
import getopt
import sys
import pax
import portage


def get_objects():

    vardb = portage.db[portage.root]["vartree"].dbapi

    objects = []

    for pkg in vardb.cpv_all():
        needed = vardb.aux_get(pkg, ['NEEDED.ELF.2'])[0].strip()
        if not needed:  # Some packages have no NEEDED.ELF.2
            continue
        for line in re.split('\n', needed):
            link = re.split(';', line)
            objects.append(link[1])  # link[1] is the ELF object

    return objects


def run_usage():
    print('Package Name : elfix')
    print('Bug Reports  : http://bugs.gentoo.org/')
    print('Program Name : migrate')
    print('Description  : Migrate PT_PAX to XATTR_PAX Flags on all system ELF objects')
    print('')
    print('Usage        : migrate -v        print out all system ELF objects')
    print('             : migrate -m [-v]   migrate flags on all system ELF objects')
    print('             : migrate -d [-v]   delete XATTR_PAX on all system ELF objects')
    print('             : migrate [-h]      print out this help')
    print('             : -v                be verbose when migrating')
    print('')


def main():
    # Are we root?
    uid = os.getuid()
    if uid != 0:
        print('This program must be run as root')
        sys.exit(1)

    try:
        opts, args = getopt.getopt(sys.argv[1:], 'vmdh')
    except getopt.GetoptError as err:
        print(str(err))  # will print something like 'option -a not recognized'
        run_usage()
        sys.exit(1)

    verbose = False
    do_migration = False
    do_deleteall = False
    do_usage = False

    opt_count = 0

    for o, a in opts:
        if o == '-v':
            verbose = True
            opt_count += 1
        elif o == '-m':
            do_migration = True
            opt_count += 1
        elif o == '-d':
            do_deleteall = True
            opt_count += 1
        elif o == '-h':
            do_usage = True
            opt_count += 1
        else:
            print('Option included in getopt but not handled here!')
            print('Please file a bug')
            sys.exit(1)

    if do_usage:
        run_usage()
        sys.exit(0)

    if opt_count == 0 or opt_count > 2 or (do_migration and do_deleteall):
        run_usage()
        sys.exit(1)

    # Do we have XATTR_PAX support?
    if do_migration or do_deleteall:
        try:
            from pax import deletextpax
        except ImportError:
            print('ERROR: Python module pax.so was compiled without XATTR_PAX support, '
                  'cannot migrate or delete XATTR_PAX')
            sys.exit(1)

    objects = get_objects()

    fail = []
    none = []

    for elf in objects:
        try:
            flags = pax.getflags(elf)[0]
            if flags:
                if verbose:
                    print("%s %s" % (flags, elf))
            else:
                none.append(elf)
                if verbose:
                    print("NONE: %s" % elf)

            if do_migration:
                flags = re.sub('-', '', flags)
                if flags == 'e':
                    continue  # Don't create XATTR_PAX for default
                pax.setstrflags(elf, flags)

            if do_deleteall:
                pax.deletextpax(elf)

        # We should never get here, because you can
        # always getflags() via pax.so since you can
        # read PT_PAX even from a busy text file, and
        # you can always set the pax flags with pax.so
        # even on a busy text file because it will skip
        # setting PT_PAX and only set the XATTR_PAX
        except pax.PaxError:
            if uid == 0:
                fail.append(elf)
                if verbose:
                    print("FAIL: %s" % elf)

    if verbose:
        if fail:
            print('\n')
            print("ELF executables for which the migration failed:")
            for elf in fail:
                print("\t%s" % elf)
        if none:
            print('\n')
            print("ELF executables lacking PT_PAX:")
            for elf in none:
                print("\t%s" % elf)

if __name__ == '__main__':
    main()