#!/usr/bin/env python2

import sys
import subprocess
import re
import os.path

# Get the version-release of the component for which rpm provides a
# live patch such that the live patch is equivalent to that component
# at version-release.
def get_complete_version(rpm):
    out = subprocess.check_output(["rpm", "-qp", "--provides", rpm])
    for line in out.split("\n"):
        match = re.match("livepatch\((.*)\)$", line.strip())
        if match:
            return match.group(1)

# Get the build-id of the binary which this live patch is meant to be
# applied on.
def get_target_buildid(rpm):
    out = subprocess.check_output(["rpm", "-qp", "--provides", rpm])
    for line in out.split("\n"):
        match = re.match("livepatch-target-buildid\((.*)\)$", line.strip())
        if match:
            return match.group(1)

# Returns a fragment of script which checks if the given complete live patch
# would apply on a running system.
def format_complete_script(component, buildid, version, release, count):
    start = 'if' if count == 0 else 'elif'

    return '''%s livepatch %s-check %s %s %s; then
    num_complete=$((num_complete+1))
''' % (start, component, version, release, buildid)

# Returns a fragment of script which checks if the given live patch would
# apply on a running system.
def format_other_script(component, buildid, version, release):
    return 'livepatch %s-check %s %s %s && num_other=$((num_other+1))\n' % (component, version, release, buildid)

def fail(s):
    print >> sys.stderr, 'ERROR: ' + s
    exit(1)

components = [{'component': 'xen',
               'pkg': 'xen-hypervisor',
               'lp-pkg': 'xen-livepatch'},
              {'component': 'xen',
               'pkg': 'xen-hypervisor',
               'lp-pkg': 'xen-ely-livepatch'},
              {'component': 'kernel',
               'pkg': 'kernel',
               'lp-pkg': 'kernel-livepatch'}]

rpms = sys.argv[1:]

complete_script = ''
other_script = ''
mainpkgs = set()

for comp in components:
    num_complete_found = 0
    mainpkg = ''

    # Look for the version-release of the main package contained in this update
    # associated with component, e.g. xen-hypervisor for xen.
    for rpm in rpms:
        match = re.match("%s-([0-9].*)\.x86_64\.rpm$" % comp['pkg'], os.path.basename(rpm))
        if match:
            if mainpkg:
                fail('Found multiple main packages for component ' + comp['component'])
            mainpkg = "%s-%s" % (comp['component'], match.group(1))

    if mainpkg:
        mainpkgs.add(mainpkg)

    # Look for live patches for this component and separate into complete live patches
    # and other live patches.
    for rpm in rpms:
        match = re.match("%s_.*([0-9]+)-([0-9]+)\.x86_64\.rpm$" % comp['lp-pkg'], os.path.basename(rpm))
        if match:
            if not mainpkg:
                fail('Found live patches for %s with no main package' % comp['component'])

            version = match.group(1)
            release = match.group(2)
            complete_version = get_complete_version(rpm)
            buildid = get_target_buildid(rpm)
            if complete_version and complete_version == mainpkg:
                # The live patch is complete for the version of <component> inside
                # this update.
                complete_script += format_complete_script(comp['component'], buildid, version, release, num_complete_found)
                num_complete_found += 1
            else:
                other_script += format_other_script(comp['component'], buildid, version, release)

    if num_complete_found > 0:
        complete_script += 'fi\n'

# Output a POSIX shell script which:
# - Checks which complete live patches are applicable.
# - If there is a complete, applicable live patch for every
#   component in this update, returns COMPLETE.
# - Checks which other live patches are applicable.
# - If any are applicable, returns INCOMPLETE.
# - Else returns NOT_APPLICABLE.

print '''#!/bin/sh

num_complete=0
num_other=0\n'''
print complete_script
print '''if [ "$num_complete" -eq %d ]; then
    # Every live patchable component in this update has a complete live patch
    # that can be applied.
    echo PATCH_PRECHECK_LIVEPATCH_COMPLETE
    exit 0
fi\n''' % len(mainpkgs)
print other_script
print '''if [ "$num_other" -gt 0 ]; then
    # At least one live patch can be applied.
    echo PATCH_PRECHECK_LIVEPATCH_INCOMPLETE
else
    # There are no live patches that can be applied.
    echo PATCH_PRECHECK_LIVEPATCH_NOT_APPLICABLE
fi

exit 0
'''
