#!/bin/bash
# vim: set et sw=4 sts=4 tw=80:
# Copyright 2007-2009 Gentoo Foundation
# Distributed under the terms of the GNU General Public License v2

# A bit of hackery to update everything that is humanly possible
# that maybe related to an older version of Python. This script can
# be run as many times as you like.
#
# OLD_PYTHON_VERSIONS = Old Python versions we are upgrading from.
# NEW_PYTHON_VERSION  = New Python version we are upgrading to.
# PKGS_EXCEPTIONS     = Packages that should NOT be re-emerged for any reason.
# PKGS_MANUAL         = Packages that should be re-emerged even if they don't fit
#                       the criteria (eg. ones that have Python compiled statically).
#
# Runtime Variables:
#
# PKGS_TO_REMERGE     = List of packages we deem to need re-emerging.
# PKGS_OK             = List of packages that should be merged without any problems.
# PKGS_MISSING        = List of packages that are installed, but cannot be merged,
#                       because they have been pruned from portage.

VERSION="0.8"
OLD_PYTHON_VERSIONS=""
NEW_PYTHON_VERSION="$(/usr/bin/python -c 'from sys import version_info; print(".".join([str(x) for x in version_info[:2]]))')"

PKGS_EXCEPTIONS="dev-lang/python sys-apps/portage"
PKGS_MANUAL="app-office/gnumeric app-office/dia dev-libs/boost x11-libs/vte"

PRETEND=0
REINSTALL_IDENTICAL_VERSIONS=0
VERBOSE=0
PKGS_TO_REMERGE=""
PKGS_COUNT_REMERGE=0
PYTHON_VERSIONS="2.1 2.2 2.3 2.4 2.5 2.6 2.7 3.0 3.1 3.2 3.3"

SUPPORTED_PMS="portage pkgcore paludis"
PMS_COMMAND=( "emerge" "pmerge" "paludis" )
PMS_OPTIONS=( "-Dv1 --keep-going" "-Do" "-i1" )
PMS_INDEX=0
CUSTOM_PMS_COMMAND=""
ADDITIONAL_OPTIONS=""

# Checks
CHECK_ECLASS=0
CHECK_MANUAL=1
CHECK_PYLIBDIR=1
CHECK_PYTHON_ABIS=1
CHECK_SONAME=1
CHECK_ECLASS_NEED_REBUILD=1

# load the gentoo-style info macros, but hack to get around
# it thinking this is an rc script
EBUILD="1"
source /etc/init.d/functions.sh

# portage variables
PKG_DBDIR=/var/db/pkg

# usage()
# display usage
usage() {
    cat <<EOF_USAGE
${0##*/} -- Find & rebuild packages broken due to a Python upgrade

Usage: python-updater [OPTION]

Options:
    -h, --help      Print usage
    -V, --version   Print version
    -p, --pretend   Pretend (don't do anything)
    -v, --verbose   Increase verbosity (may be specified multiple times)
    --reinstall-identical-versions
                    Reinstall identical versions of packages
    -P PM, --package-manager PM
                    Use package manager PM, where PM can be one of:
$(for p in ${SUPPORTED_PMS}; do
echo -ne $'\t\t    '\* ${p}
if [[ ${p} == portage ]]; then
    echo ' (Default)'
else
    echo
fi
done)
    -c CMD, --command CMD
                    Pipe found packages to command CMD instead of invoking package
                    manager. Only for debug and script use.
    --package-manager-command CMD
                    Call CMD instead of the default command for the specified
                    package manager.
    -eCHECK --enable-CHECK
                    Enable CHECK where CHECK can be one of:
                    * eclass       (Disabled by default)
                    * pylibdir     (Enabled by default)
                    * PYTHON_ABIS  (Enabled by default)
                    * soname       (Enabled by default)
                    * manual       (Enabled by default)
                    * need_rebuild (Enabled by default)
    -dCHECK --disable-CHECK
                    Disable CHECK. See --enable option for a list of checks
    -- OPTIONS      Pass additional options to package manager

See CHECKS section in the manpage for explanations about checks and
EXAMPLES section to learn how to use them.
EOF_USAGE
}

# veinfo(verbosity, message)
# einfo message if VERBOSE is bigger than verbosity.
veinfo() {
    if [[ VERBOSE -ge $1 ]]; then
        shift
        einfo "$@"
    fi
}

# die(message)
# Print error message and exit.
die() {
    eerror "$@"
    exit 1
}

# has_version(cpv)
has_version() {
    if [[ "${PMS_COMMAND[${PMS_INDEX}]}" == "emerge" ]]; then
        /usr/bin/portageq has_version / "$1"
    elif [[ "${PMS_COMMAND[${PMS_INDEX}]}" == "pmerge" ]]; then
        pquery --vdb "$1"
    else
        paludis --has-version "$1"
    fi
}

# get_vdb_variable(variable, environment_file)
get_vdb_variable() {
    local variable="$1"
    local environment_file="$2"

    unset "${variable}"
    local output
    output="$(bzip2 -cd "${environment_file}" | grep -E "^declare (-[[:alnum:]-]+ )?${variable}=")"
    if [[ -z "${output}" ]]; then
        output="$(bzip2 -cd "${environment_file}" | grep -E "^(declare (-[[:alnum:]-]+ )?)?${variable}=")"
    fi
    if [[ -n "${output}" ]]; then
        output="$(echo "${output}" | sed -e "s/^\(declare \(-[[:alnum:]-]\+ \)\?\)\?${variable}=\(\"\)\?//;s/\(\"\)\?$//")"
        if [[ "${output}" == "'"*"'" ]]; then
            output="${output:1:$((${#output}-2))}"
        fi
        eval "${variable}=\"\${output}\""
        return 0
    else
        return 1
    fi
}

USE_PYTHON=""

# get_USE_PYTHON(CATEGORY, PN, SLOT)
get_USE_PYTHON() {
    if [[ "${PMS_COMMAND[${PMS_INDEX}]}" =~ ^(emerge|pmerge)$ ]]; then
        if [[ -z "${USE_PYTHON}" ]]; then
            USE_PYTHON="$(/usr/bin/portageq envvar USE_PYTHON)"
        fi
    else
        USE_PYTHON="$(paludis --environment-variable "$1/$2:$3::/?" USE_PYTHON)"
    fi
    if [[ -z "${USE_PYTHON}" ]]; then
        USE_PYTHON="${PYTHON2_VERSION} ${PYTHON3_VERSION}"
        USE_PYTHON="${USE_PYTHON# }"
        USE_PYTHON="${USE_PYTHON% }"
    fi
    echo "${USE_PYTHON}"
}

# get_ebuild_path(CATEGORY, PN, SLOT)
get_ebuild_path() {
    # These manual calls to Portage API are 2 times faster than 'portageq get_repo_path / "$(portageq metadata / ebuild "$(portageq best_visible / "$1")" repository)"'.
    /usr/bin/python -c "
import sys

if '' in sys.path:
    sys.path.remove('')

import portage

dbapi = portage.db['/']['porttree'].dbapi
pkg_list = dbapi.match('$1/$2:$3')
if pkg_list:
    best_visible = portage.best(pkg_list)
    if best_visible:
        repository = dbapi.aux_get(best_visible, ['repository'])
        if repository:
            repository_path = dbapi.getRepositoryPath(repository[0])
            if repository_path:
                print(repository_path + '/$1/$2/' + best_visible.split('/')[1] + '.ebuild')"
}

# get_RESTRICT_PYTHON_ABIS(CATEGORY, PN, SLOT)
get_RESTRICT_PYTHON_ABIS() {
    if [[ "${PMS_COMMAND[${PMS_INDEX}]}" =~ ^(emerge|pmerge)$ ]]; then
        local ebuild_path="$(get_ebuild_path "$1" "$2" "$3")"
        if [[ -f "${ebuild_path}" ]]; then
            grep "^RESTRICT_PYTHON_ABIS=" "${ebuild_path}" | sed -e "s/RESTRICT_PYTHON_ABIS=[\"']\(.*\)[\"']/\1/"
        fi
    else
        paludis --environment-variable "$1/$2:$3::/?" RESTRICT_PYTHON_ABIS
    fi
}

# Respect PYUPDATER_OPTIONS
if [[ -n "${PYUPDATER_OPTIONS}" ]]; then
    set -- ${PYUPDATER_OPTIONS} "$@"
fi

# Respect PACKAGE_MANAGER environment variable
i=0
for pm in ${SUPPORTED_PMS}; do
    if [[ "${pm}" == "${PACKAGE_MANAGER}" ]]; then
        PMS_INDEX="${i}"
        break
    fi
    ((i++))
done

# Command Line Parsing
while [[ -n "$1" ]]; do
    case "$1" in
        -h|--help)
            usage
            exit 0
            ;;
        -V|--version)
            echo "${VERSION}"
            exit 0
            ;;
        -p|--pretend)
            PRETEND=1
            ;;
        -v|--verbose)
            VERBOSE=$(( $VERBOSE + 1 ))
            ;;
        --reinstall-identical-versions)
            REINSTALL_IDENTICAL_VERSIONS=1
            ;;
        -P|--package-manager)
            shift
            PACKAGE_MANAGER="$1"
            case "${PACKAGE_MANAGER}" in
                portage|pkgcore|paludis)
                    ;;
                *)
                    echo "Unrecognized package manager selected. Please select between ${SUPPORTED_PMS}"
                    exit 1
                    ;;
            esac

            # PMS_INDEX is used to select the right commands and options for the selected package manager
            PMS_INDEX=0
            for PM in ${SUPPORTED_PMS}; do
                [[ ${PM} == ${PACKAGE_MANAGER} ]] && break
                PMS_INDEX=$((${PMS_INDEX} + 1))
            done
            ;;
        --package-manager-command)
            shift
            CUSTOM_PMS_COMMAND="$1"
            ;;
        -c|--command)
            shift
            PIPE_COMMAND="$1"
            ;;
        -ee*|--enable-e*)
            CHECK_ECLASS=1
            ;;
        -de*|--disable-e*)
            CHECK_ECLASS=0
            ;;
        -em*|--enable-m*)
            CHECK_MANUAL=1
            ;;
        -dm*|--disable-m*)
            CHECK_MANUAL=0
            ;;
        -eP*|--enable-P*)
            CHECK_PYTHON_ABIS=1
            ;;
        -dP*|--disable-P*)
            CHECK_PYTHON_ABIS=0
            ;;
        -ep*|--enable-p*)
            CHECK_PYLIBDIR=1
            ;;
        -dp*|--disable-p*)
            CHECK_PYLIBDIR=0
            ;;
        -es*|--enable-s*)
            CHECK_SONAME=1
            ;;
        -ds*|--disable-s*)
            CHECK_SONAME=0
            ;;
        -en*|--enable-n*)
            CHECK_ECLASS_NEED_REBUILD=1
            ;;
        -dn*|--disable-n*)
            CHECK_ECLASS_NEED_REBUILD=0
            ;;
        --)
            shift
            ADDITIONAL_OPTIONS="$*"
            break
            ;;
        *)
            usage
            echo "Unrecognized option '$1'"
            exit 1
            ;;
    esac
    shift
done

PYTHON2_VERSION=""
PYTHON3_VERSION=""
if has_version "=dev-lang/python-2*"; then
    if [[ "$(readlink /usr/bin/python2)" != "python2."* ]]; then
        die "'/usr/bin/python2' is not valid symlink"
    fi
    PYTHON2_VERSION="$(/usr/bin/python2 -c 'from sys import version_info; print(".".join([str(x) for x in version_info[:2]]))')"
fi
if has_version "=dev-lang/python-3*"; then
    if [[ "$(readlink /usr/bin/python3)" != "python3."* ]]; then
        die "'/usr/bin/python3' is not valid symlink"
    fi
    PYTHON3_VERSION="$(/usr/bin/python3 -c 'from sys import version_info; print(".".join([str(x) for x in version_info[:2]]))')"
fi
if [[ -z "${PYTHON2_VERSION}" && -z "${PYTHON3_VERSION}" ]]; then
    die "Python 2 and Python 3 not installed"
fi

for python_version in ${PYTHON_VERSIONS}; do
    if [[ "${python_version}" != "${NEW_PYTHON_VERSION}" ]]; then
        OLD_PYTHON_VERSIONS+="${OLD_PYTHON_VERSIONS:+ }${python_version}"
    fi
done
OLD_PYTHON_VERSIONS_REGEX="(${OLD_PYTHON_VERSIONS// /|})"
OLD_PYTHON_VERSIONS_REGEX="${OLD_PYTHON_VERSIONS_REGEX//./\.}"

einfo "Starting Python Updater from ${OLD_PYTHON_VERSIONS// /, } to ${NEW_PYTHON_VERSION}:"

if [[ CHECK_SONAME -ne 0 ]]; then
    if ! type -P scanelf >/dev/null 2>&1; then
        ewarn "scanelf not found!"
        ewarn "check soname is disabled."
        CHECK_SONAME=0
    else
        veinfo 1 'check "soname" enabled.'
    fi
else
    veinfo 1 'check "soname" disabled.'
fi
[[ CHECK_PYLIBDIR -ne 0 ]] \
    && veinfo 1 'check "pylibdir" enabled.' \
    || veinfo 1 'check "pylibdir" disabled.'
[[ CHECK_ECLASS -ne 0 ]] \
    && veinfo 1 'check "eclass" enabled.' \
    || veinfo 1 'check "eclass" disabled.'
[[ CHECK_MANUAL -ne 0 ]] \
    && veinfo 1 'check "manual" enabled.' \
    || veinfo 1 'check "manual" disabled.'
[[ CHECK_ECLASS_NEED_REBUILD -ne 0 ]] \
    && veinfo 1 'check "need_rebuild" enabled.' \
    || veinfo 1 'check "need_rebuild" disabled.'

# Iterate through the contents of all the installed packages.
# ${PKG_DBDIR} must be followed by '/' to avoid problems when ${PKG_DBDIR} is a symlink.
for content in $(find ${PKG_DBDIR}/ -name CONTENTS | sort); do
    environment_file="${content/CONTENTS/environment.bz2}"

    # Extract some variables. PYTHON_ABIS is optional.
    get_vdb_variable CATEGORY "${environment_file}" || die "CATEGORY missing"
    get_vdb_variable PN "${environment_file}" || die "PN missing"
    get_vdb_variable PF "${environment_file}" || die "PF missing"
    get_vdb_variable SLOT "${environment_file}" || die "SLOT missing"
    get_vdb_variable PYTHON_ABIS "${environment_file}"
    CATPKG="${CATEGORY}/${PN}"
    CATPKGVER="${CATEGORY}/${PF}"
    veinfo 2 "Checking ${CATEGORY}/${PF}:${SLOT}"

    # Exclude packages that are exceptions, like Portage and Python itself.
    exception=0
    for exp in ${PKGS_EXCEPTIONS}; do
        if [[ -z "${CATPKG##${exp}}" ]]; then
            veinfo 2 "Skipping ${CATPKG}, reason: exception"
            exception=1
            break;
        fi
    done

    [[ ${exception} == 1 ]] && continue

    # Check if package is in PKGS_MANUAL
    if [[ CHECK_MANUAL -ne 0 ]]; then
        for pkg in ${PKGS_MANUAL}; do
            if [[ -z "${CATPKG##${pkg}}" ]]; then
                exception=2
                break;
            fi
        done
    fi

    # Replace SLOT by version number when REINSTALL_IDENTICAL_VERSIONS == 1
    # Reinstall identical versions when SLOT doesn't exist, bug #201848
    if [[ REINSTALL_IDENTICAL_VERSIONS -eq 1 || ! -f "${content/CONTENTS/SLOT}" ]]; then
        CATPKGVER="=${CATPKGVER}"
    else
        SLOT="$(<"${content/CONTENTS/SLOT}")"
        CATPKGVER="${CATPKG}:${SLOT}"
    fi

    if [[ ${exception} = 2 ]]; then
        PKGS_TO_REMERGE+=" ${CATPKGVER}"
        eindent
        einfo "Adding to list: ${CATPKGVER}"
        eindent
        einfo "check: manual [Added to list manually, see CHECKS in manpage for more information.]"
        eoutdent && eoutdent
        continue
    fi

    if [[ CHECK_PYTHON_ABIS -ne 0 ]]; then
        if [[ -n "${PYTHON_ABIS}" ]]; then
            new_PYTHON_ABIS=""
            RESTRICT_PYTHON_ABIS="$(get_RESTRICT_PYTHON_ABIS "${CATEGORY}" "${PN}" "${SLOT}")"
            for PYTHON_ABI in $(get_USE_PYTHON "${CATEGORY}" "${PN}" "${SLOT}"); do
                support_ABI="1"
                for restricted_ABI in ${RESTRICT_PYTHON_ABIS}; do
                    if [[ "${PYTHON_ABI}" == ${restricted_ABI} ]]; then
                        support_ABI="0"
                        break
                    fi
                done
                [[ "${support_ABI}" == "1" ]] && new_PYTHON_ABIS+="${new_PYTHON_ABIS:+ }${PYTHON_ABI}"
            done
            if [[ "${PYTHON_ABIS}" != "${new_PYTHON_ABIS}" ]]; then
                PKGS_TO_REMERGE+=" ${CATPKGVER}"
                eindent
                einfo "Adding to list: ${CATPKGVER}"
                eindent
                veinfo 1 "check: PYTHON_ABIS [ Previous Python ABIs ${PYTHON_ABIS}, new Python ABIs ${new_PYTHON_ABIS} ]"
                eoutdent && eoutdent
                continue
            fi
            # Don't run other checks if PYTHON_ABIS check has been run.
            continue
        fi
    fi

    if [[ CHECK_ECLASS_NEED_REBUILD -ne 0 ]]; then
        get_vdb_variable PYTHON_NEED_REBUILD "${environment_file}"
        if echo "${PYTHON_NEED_REBUILD}" | grep -qE "${OLD_PYTHON_VERSIONS_REGEX}"; then
            PKGS_TO_REMERGE+=" ${CATPKGVER}"
            eindent
            einfo "Adding to list: ${CATPKGVER}"
            eindent
            veinfo 1 "check: need_rebuild [ Ebuild set PYTHON_NEED_REBUILD=${PYTHON_NEED_REBUILD} ]"
            eoutdent && eoutdent
            continue
        fi
    fi

    if [[ CHECK_PYLIBDIR -ne 0 ]]; then
        # Search for possible old Python dirs in CONTENTS
        # /usr/include/python$old
        # /usr/lib/python$old
        # /usr/lib32/python$old
        # /usr/lib64/python$old
        if grep -qE "/usr/(include|lib(32|64)?)/python${OLD_PYTHON_VERSIONS_REGEX}" "${content}"; then
            PKGS_TO_REMERGE+=" ${CATPKGVER}"
            eindent
            einfo "Adding to list: ${CATPKGVER}"
            eindent
            veinfo 1 "check: pylibdir [ Installed file under old Python include/library directory ]"
            eoutdent && eoutdent
            continue
        fi
    fi

    if [[ CHECK_SONAME -ne 0 ]]; then
        broken_libs=""
        for python_version in ${OLD_PYTHON_VERSIONS}; do
            soname="$(readlink -n /usr/lib/libpython${python_version}.so)"
            if [[ -z "${soname}" ]]; then
                # Use safe, default value.
                soname="libpython${python_version}.so.1.0"
            fi
            broken_libs+="$(scanelf -qBN "${soname}" < <(grep -E "^obj" "${content}" | cut -d' ' -f2))"
        done
        if [[ -n "${broken_libs}" ]]; then
            PKGS_TO_REMERGE+=" ${CATPKGVER}"
            eindent
            einfo "Adding to list: ${CATPKGVER}"
            eindent
            veinfo 1 "check: soname [ Libraries linked to old libpython found:"
	    eindent
	    old_IFS="${IFS}"
	    IFS=$'\n'
            for broken_lib in ${broken_libs}; do
                veinfo 1 "${broken_lib}"
            done
	    IFS="${old_IFS}"
	    eoutdent
            veinfo 1 "]"
            eoutdent && eoutdent
        fi
    fi

    if [[ CHECK_ECLASS -ne 0 ]]; then
        get_vdb_variable PYVER "${environment_file}"
        if echo "${PYVER}" | grep -qE "${OLD_PYTHON_VERSIONS_REGEX}"; then
            PKGS_TO_REMERGE+=" ${CATPKGVER}"
            eindent
            einfo "Adding to list: ${CATPKGVER}"
            eindent
            veinfo 1 "check: eclass [ Ebuild set PYVER=${PYVER} ]"
            eoutdent && eoutdent
            continue
        fi
    fi
done

# Pipe to command if we have one
if [[ -n "${PIPE_COMMAND}" ]]; then
    echo "${PKGS_TO_REMERGE}" | ${PIPE_COMMAND}
    exit $?
fi

if [[ "${PMS_COMMAND[${PMS_INDEX}]}" == "emerge" ]] ; then
    # Filter out --getbinpkg, --getbinpkgonly, --usepkg and --usepkgonly options in EMERGE_DEFAULT_OPTS environment variable
    emerge_default_opts=""
    for option in $(/usr/bin/portageq envvar EMERGE_DEFAULT_OPTS); do
        if [[ "${option}" == -[[:alnum:]]* ]]; then
            [[ "${option//[gGkK]/}" != "-" ]] && emerge_default_opts+=" ${option//[gGkK]/}"
        elif [[ "${option}" != "--getbinpkg" && "${option}" != "--getbinpkgonly" && "${option}" != "--usepkg" && "${option}" != "--usepkgonly" ]]; then
            emerge_default_opts+=" ${option}"
        fi
    done
    export EMERGE_DEFAULT_OPTS="${emerge_default_opts# }"
fi

# Only pretending?
[[ PRETEND -eq 1 ]] && PMS_OPTIONS[${PMS_INDEX}]="${PMS_OPTIONS[${PMS_INDEX}]} -p"

# (Pretend to) reinstall packages
if [[ -n "${PKGS_TO_REMERGE}" ]]; then
    pmscmd="${CUSTOM_PMS_COMMAND}"
    [[ -z "${pmscmd}" ]] && pmscmd="${PMS_COMMAND[${PMS_INDEX}]}"
    cmd="${pmscmd} ${PMS_OPTIONS[${PMS_INDEX}]} ${PKGS_TO_REMERGE} ${ADDITIONAL_OPTIONS}"
    einfo ${cmd}
    ${cmd}
else
    einfo "No packages need to be reinstalled."
fi