#!/usr/bin/env python
# Copyright (c) 2020-2021 by Cisco Systems, Inc.
 
import os
import sys
import fcntl
import argparse
import subprocess
import re
try:
    import provides_hint
    RPM_TOOL = "rpm"    
except:
    if os.path.exists ("/usr/bin/rpm.orig"):
        RPM_TOOL = "rpm.orig"
    else:
        RPM_TOOL = "rpm"
    pass

#import pdb

YML_LOCATION  = "/install/instdb/local/"
YML_FILE      = "iosxr_image_mdata.yml"

'''
Current arbitrary tags in code.
Parentversion, PIPD and skiprelease is not populated
as part of XR packaging, but used in install code.
default all to none for rpm4 output.
'''
custom_mdata = {
                "SUPPCARDS" : '(none)', 
                "VMTYPE": '(none)',
                "CISCOHW" : '(none)', 
                "PACKAGETYPE" : '(none)',
                "PKGTYPE" : '(none)', 
                "PACKAGEPRESENCE" : '(none)', 
                "RESTARTTYPE" : '(none)',
                "INSTALLMETHOD" : '(none)', 
                "XRVERSION" : '(none)', 
                "XRRELEASE" : '(none)',
                "PARENTVERSION" : '(none)', 
                "PIPD" : '(none)', 
                "SKIPRELEASE" : '(none)', 
                "CARDTYPE" : '(none)'
               }
def run_cmd (cmd):
    process = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE, shell=True)
    out, error = process.communicate()
    sprc = process.returncode
    if sprc is None or sprc != 0:
        out = error
        raise RuntimeError("Error CMD=%s returned --->%s" % (cmd, out))
    else:
        ''' rpm returns exit code 0 even when error is populated '''
        if not out:
            out = error
            raise RuntimeError("Error CMD=%s returned --->%s" % (cmd, out))
    return out.strip()

''' 
Query GROUP tag for custom metadata in package.
If result has tag like suppcards, GROUP tag holds custom metadata.
else, continue with original query tag.
In case, we have mix of packages in query, raise an exception.
'''
def query_custom_tag (opts, is_pkg):
    query_grp_tag = False
    rpm_supports_cust_tags = False
    groupdata = {}
    cust_tag_support_checked = False
    for pkg in opts:
        if is_pkg and not os.path.isfile (pkg):
            continue
        try:
            stroutput = None
            if is_pkg :
                rpmcmd = "{} -qp --qf {} {}".format (
                         RPM_TOOL, "\'[%{GROUP}\\n]\'", pkg
                         )
            else :
                rpmcmd = "{} -q --qf {} {}".format (
                         RPM_TOOL, "\'[%{GROUP}\\n]\'", pkg
                         )

            stroutput = run_cmd (rpmcmd)
            if 'SUPPCARDS' in stroutput.upper() or 'XRRELEASE' in stroutput.upper():
                query_grp_tag = True
                groupdata[pkg] = stroutput
            else:
                # If GROUP field doesnt have the custom tags, check if rpm supports
                # custom tags (this will be applicable to platforms supporting rpm5)
                if not cust_tag_support_checked:
                    try:
                        cust_tag_support_checked = True 
                        stroutput = ""
                        if is_pkg:
                            rpmcmd = "{} -qp --qf {} {}".format (
                                     RPM_TOOL, "\'[%{XRRELEASE} %{VMTYPE} %{SUPPCARDS}\\n]\'", pkg
                                     )
                        else: 
                            rpmcmd = "{} -q --qf {} {}".format (
                                     RPM_TOOL, "\'[%{XRRELEASE} %{VMTYPE} %{SUPPCARDS}\\n]\'", pkg
                                     )
                        stroutput = run_cmd (rpmcmd)
                        if stroutput:
                            rpm_supports_cust_tags = True
                    except:
                        pass 
                if query_grp_tag:
                    raise Exception ("Mixed packages in input.")
                    query_grp_tag = False
        except:
            raise
    return query_grp_tag, groupdata, rpm_supports_cust_tags

'''
Parse input for query string, we are only interested in queries
for custom tags.
'''
def parseCli ():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-q", 
        dest="query",
        action='store_true',
        required=False)

    parser.add_argument(
        "--qf",
        "--queryformat",
        dest="querystring",
        required=False)
 
    parser.add_argument(
        "-g",
        dest="querygroup",
        action='store_true',
        required=False)

    parser.add_argument(
        "-p", 
        dest="querypkg", 
        action='store_true',
        required=False)

    parser.add_argument(
        "--new", 
        dest="newfmt",
        action='store_true',
        required=False)

    parser.print_usage = None

    namespace, opts = parser.parse_known_args ()
    return namespace, opts
 
def get_platform (): 
    BOOTSTRAP_FILE='/etc/rc.d/init.d/calvados_bootstrap.cfg'
    PLATFORM_STR1='PLATFORM_EXT'
    PLATFORM_STR2='PLATFORM'
    platform = None

    with open(BOOTSTRAP_FILE) as bootFile:
        for line in bootFile:
            line = line.split('#')[0].strip()
            if not line:
                continue
            (key, value) = line.split('=')
            if key == PLATFORM_STR1:
                platform = value
            elif key == PLATFORM_STR2:
                platform = value
    return platform

if __name__ == "__main__":
    querystr = None
    opts = None
    namespace = None

    try:
        namespace, opts = parseCli ()
    except:
        pass

    if namespace and namespace.querygroup:
        qa_output = ''
        empty_msg = ''
        platform = get_platform ()
        if platform:
            if "sysadmin" in opts:
                rpmcmd = {
                    "{} -qa --qf {} | grep {}-sysadmin".format (
                        RPM_TOOL,
                        "\'[%{NAME}-%{VERSION}-%{RELEASE}.%{ARCH} %{BUILDTIME}\\n]\'",
                        platform
                    )
                }
                empty_msg = "group sysadmin does not contain any packages"
            else:
                rpmcmd = {
                    "{} -qa --qf {} | grep {} | grep -v {}-sysadmin".format (
                        RPM_TOOL,
                        "\'[%{NAME}-%{VERSION}-%{RELEASE}.%{ARCH} %{BUILDTIME}\\n]\'",
                        platform,
                        platform
                    )
                }
                empty_msg = "group IOS-XR does not contain any packages"
            try:
                cmdoutput = run_cmd (rpmcmd)
                it = iter (cmdoutput.split())
                builddict = dict (zip(it, it))
                sort_dict = sorted(builddict.items(), key=lambda x: x[1], reverse = True)
                qa_output = '\n'.join([x[0] for x in sort_dict])
            except:
                print (empty_msg)

        print (qa_output)
        sys.exit (0)

    if namespace and namespace.querystring:
        querystr = namespace.querystring.lower()
    inargs = ' '.join(sys.argv[1:])
 
    '''
        Check if it is a query on a package for custom metadata.
        If so, check that the rpm being queried do hold custom tag.
        If not, get data from GROUP tag in rpm.
    '''
    query_grp_tag = False
    suppcards = ""
    install = ""
    restart_type = "" 
    pkg_type = "" 
    pkg_pres = "" 
    rel_data = ""
    xr_ver = ""
    hw = ""
    vmtype = ""
    package_type = ""
    
    if namespace and namespace.query and namespace.querystring:
        resultdict = {}
        groupdata = {}
        rpmstroutput = ''
        rpm_supports_cust_tags = False
        #pdb.set_trace()

        if any (item in querystr.upper() for item in custom_mdata.keys()):
            if namespace.newfmt:
                query_grp_tag = True
            else:
                if namespace.querypkg:
                    query_grp_tag, groupdata, rpm_supports_cust_tags = query_custom_tag (opts, True)
                else:
                    query_grp_tag, groupdata, rpm_supports_cust_tags = query_custom_tag (opts, False)

            if not query_grp_tag:
                yml_loaded = False
                if rpm_supports_cust_tags :
                    try:
                        ''' Check if rpm has needed custom tags, handover to rpm'''
                        sys.argv[0] = RPM_TOOL
                        os.execvp(RPM_TOOL, sys.argv)
                    except:
                        pass
                else:
                    try:
                        pkgdict = dict (custom_mdata)
                        for pkg in opts:
                            pkgargs = inargs
                            if namespace.querypkg and not os.path.isfile (pkg):
                                continue
                            # Check if platform name is part of rpm name
                            # If not, it is a TP smu
                            rpmcmd = "{} -qa | grep sysadmin".format (RPM_TOOL)
                            qa_output = run_cmd (rpmcmd)
                            rpm_ = qa_output.split('\n', 1)[0]
                            platform = rpm_.split('-', 1) [0]
                            #platform = "ncs5500"
                            basename = os.path.basename(pkg)
                            if platform not in basename:
                                if '.xr.' in basename:
                                    vmtype = "xr"
                                elif '.admin.' in basename:
                                    vmtype = "calvados"
                                elif '.host.' in basename:
                                    vmtype = "host"
                            rpm_dict = {"VMTYPE" : vmtype}
       
                            #Skip reading custom tag data from yml for tp rpms
                            if platform in basename:
                                if not yml_loaded:
                                    p_name = os.path.basename(pkg)
                                    rel_ss = re.split(r'(.*)-(.*)', p_name)[2]
                                    rel  = rel_ss.split(".",1)[0]
                                    #yml_name = './' + YML_FILE + '_' + rel
                                    yml_name = YML_LOCATION + YML_FILE + '_' + rel
                                    if os.path.exists(yml_name):
                                        hints = provides_hint.PkgHintsMdata(yml_name)
                                        yml_loaded = True
                                pkg_inst_dict_list = hints.get_mdata_for_key('package_instances')
                                for x in opts:
                                    if pkg == x:
                                        continue
                                    elif namespace.querypkg and not os.path.isfile (x):
                                        continue
                                    pkgargs = pkgargs.replace (x, '')
                                pkg_inst_name = re.split(r'-\d.\d.\d.*', basename)[0]
                                if "sysadmin" in pkg:
                                    pkginst_vm = pkg_inst_dict_list['pkginst_cal']
                                    pkg_substr = pkg_inst_name.split("-", 2)[2]
                                    inst_name = pkginst_vm[pkg_substr]
                                    suppcards = inst_name['cards']
                                    hw = inst_name['ciscohw']
                                    install = inst_name['install']
                                    package_type = inst_name['packagetype']
                                    pkg_pres = inst_name['pkgpresence']
                                    pkg_type = inst_name['pkgtype']
                                    package_type = inst_name['packagetype']
                                    restart_type = inst_name['restart']
                                    vmtype = inst_name['vmtype']
                                    xr_ver = inst_name['xrversion']
                                    rel_substr = re.split(r'(.*)-(.*)', basename)[2]
                                    rel_data  = rel_substr.split(".",1)[0] 
                                else:
                                    if "package_version_lock" in basename:
                                        suppcards = "rp,lc,sc,xc,fc,cc,sw,all"
                                    pkginst_vm = pkg_inst_dict_list['pkginst_xr']
                                    pkg_mdata_dict_list = hints.get_mdata_for_key('package_mdata')
                                    if (pkginst_vm.has_key(pkg_inst_name)): 
                                        pkginst_vm_pkg = pkginst_vm[pkg_inst_name]

                                        suppcards = pkginst_vm_pkg['cards']
                                        install   = pkginst_vm_pkg['install']
                                        restart_type = pkginst_vm_pkg['restart']
                                        pkg_pres = pkginst_vm_pkg['pkgpresence']
                                        pkg_type = pkginst_vm_pkg['pkgtype']
                                        package_type = pkginst_vm_pkg['packagetype'] 
                                        xr_ver = pkginst_vm_pkg['xrversion']
                                        hw = pkginst_vm_pkg['ciscohw']
                                        vmtype = pkginst_vm_pkg['vmtype']
                                        rel_substr = re.split(r'(.*)-(.*)', basename)[2]
                                        rel_data  = rel_substr.split(".",1)[0] 
                                    else:
                                        if (pkg_mdata_dict_list.has_key(pkg_inst_name)):
                                            pkg_mdata = pkg_mdata_dict_list[pkg_inst_name]

                                            suppcards = pkg_mdata['cards']
                                            install = pkg_mdata['install']
                                            pkg_pres = pkg_mdata['pkgpresence']
                                            restart_type = pkg_mdata['restart']

                                            pkg_data = pkg_mdata['packages']
                                            pkg_pie_name = pkg_data[0]['piename'].split('.')[0]
                                            rel_data = pkg_data[0]['release']
                                            if (pkginst_vm.has_key(pkg_pie_name)):
                                                pkginst_vm_pkg = pkginst_vm[pkg_pie_name]
                                                pkg_type = pkginst_vm_pkg['pkgtype']
                                                package_type = pkginst_vm_pkg['packagetype'] 
                                                xr_ver = pkginst_vm_pkg['xrversion']
                                                hw = pkginst_vm_pkg['ciscohw']
                                                vmtype = pkginst_vm_pkg['vmtype']

                                #Build the dictionary with the read values
                                rpm_dict = {"SUPPCARDS"       : suppcards, 
                                            "INSTALLMETHOD"   : install,
                                            "RESTARTTYPE"     : restart_type,
                                            "PACKAGETYPE"     : package_type,
                                            "PKGTYPE"         : pkg_type,
                                            "PACKAGEPRESENCE" : pkg_pres,
                                            "XRRELEASE"       : rel_data,
                                            "XRVERSION"       : xr_ver,
                                            "CISCOHW"         : hw,
                                            "VMTYPE"          : vmtype}
                                '''custom tag SUPPCARDS used to hold data with ',' as delimiter'''
                                if rpm_dict.has_key('SUPPCARDS'):
                                    rpm_dict['SUPPCARDS'] = ','.join(rpm_dict['SUPPCARDS'].split(':'))

                            # For tp smus default values are returned for custom tags (flow just comes here) 
                            pkgdict.update (rpm_dict)

                            newquerystr = namespace.querystring
                            for item in custom_mdata.keys():
                                newquerystr = re.sub ("%{{{}}}".format(item), pkgdict[item], newquerystr, flags=re.I)
                            newquerystr = newquerystr.strip('[]')
                                    
                            ''' Run the rpm query '''
                            try:
                                pkgargs = pkgargs.replace (namespace.querystring, "\'%s\'"%(newquerystr))
                                rpmcmd = "{} {}".format(RPM_TOOL, pkgargs)
                                stroutput = run_cmd(rpmcmd)
                                resultdict[pkg] = stroutput
                            except:
                                raise
                        rpmstroutput = '\n'.join(resultdict.values())
                    except:
                        pass        
            else:
                for pkg in opts:
                    pkgargs = inargs
                    if namespace.querypkg and not os.path.isfile (pkg):
                        continue
                    for x in opts:
                        if pkg == x:
                            continue
                        elif namespace.querypkg and not os.path.isfile (x):
                            continue
                        pkgargs = pkgargs.replace (x, '')
                    resultdict[pkg] = ''
                    stroutput = groupdata[pkg]
                    data = stroutput.split(',')[-1]
                    cfg = dict([(item.partition(':')[0].upper(),
                                item.partition(':')[2]) 
                                for item in data.split(';') if not 
                                item.strip().startswith('#') and item.strip()])
                    '''custom tag SUPPCARDS used to hold data with ',' as delimiter'''
                    if cfg.has_key('SUPPCARDS'):
                        cfg['SUPPCARDS'] = ','.join(cfg['SUPPCARDS'].split('-'))
                    pkgdict = dict (custom_mdata)
                    pkgdict.update (cfg)
                    ''' Edit querystr to query for standard tags'''
                    newquerystr = namespace.querystring
                    for item in custom_mdata.keys():
                        newquerystr = re.sub ("%{{{}}}".format(item), pkgdict[item], newquerystr, flags=re.I)
                    newquerystr = newquerystr.strip('[]')

                    ''' Run the rpm query '''
                    try:
                        pkgargs = pkgargs.replace (namespace.querystring, "\'%s\'"%(newquerystr))
                        rpmcmd = "{} {}".format(RPM_TOOL, pkgargs)
                        stroutput = run_cmd(rpmcmd)
                        resultdict[pkg] = stroutput
                    except:
                        raise
	        rpmstroutput = '\n'.join(resultdict.values())
        else:
            try:
                ''' 
                Query on package but not for custom tag. Hand over control 
                to rpm 
                '''
                sys.argv[0] = RPM_TOOL
                os.execvp(RPM_TOOL, sys.argv)
            except:
                pass

        print (rpmstroutput)
    else:
        ''' It is not a query on package, hand over control to rpm '''
        try:
            sys.argv[0] = RPM_TOOL
            os.execvp(RPM_TOOL, sys.argv)
        except:
            pass
    sys.exit (0)
