###############################################################################
# Local Security Check Automation Framework
#
# Authors:
# Veerendra GG <veerendragg@secpod.com>
#
# Revision 1.0
# Date: 2008/12/22
#
# Copyright:
# Copyright (c) 2009 SecPod , http://www.secpod.org
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2
# (or any later version), 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.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
###############################################################################

"""
ganerate_script: NVT Generator
Generate scripts according to the template
"""

import re
import sys
import string
import utils


class GenerateCode:
    """
    Implements generic set of functions to generate the local security check code
    """

    def generateReleaseCheck(self, release, debug=0):
        """
        Generate if open
        """

        if_open = """
if(release == "%s")
{\n""" % (release)
        return(if_open)


    def generateClose(self, debug=0):
        """
        Generates if close
        """

        if_close = """
  exit(0);
}\n"""
        return(if_close)


    def generateRPMCheck(self, parse, platform, debug = 0):
        """
        Generates RPM verifying code
        """

        try:
            rpm_code = ''

            if debug:
                print "Generating RPM Code..."

            for rpm in parse.Packages[platform]:
                release = re.findall("(?<=-|_)\d+\.\d+.*", rpm)
                if not release:
                    release = re.findall("(?<=-|_)(\d.*)", rpm)
                    if not release:
                        if debug:
                            print "Didn't find version information for RPM: ", rpm
                        continue

                release = release[0]
                package = rpm.replace(release, '')
                package = package.rstrip('-').rstrip('_')

                if not(package and release):
                    if debug:
                        print "Didn't find package and rpm information: ", package, release
                    continue

                rpm = package + '~' + release.replace("-","~")

                rpm_code += """
  if(isrpmvuln(pkg:"%s", rpm:"%s", rls:"%s"))
  {
    security_hole(0);
    exit(0);
  }\n""" % (package, rpm, platform)

            return(rpm_code)
        except Exception, msg:
            print "Exception in : GenerateCode -> generateRPMCheck() method"
            sys.exit(msg)


    def generateDebPkgCheck(self, parse, platform, debug = 0):
        """
        Generates Debian PKG verifying code
        """

        try:
            deb_pkg_code = ''

            if debug:
                print "Generating Debian Pkg Code..."

            for deb_pkg in parse.Packages[platform]:
                release = deb_pkg.split("_")
                if not release:
                  if debug:
                    print "Didn't find version information for Debian PKG: ", deb_pkg
                  continue

                package = release[0]
                version = release[1]
                package = package.rstrip('-').rstrip('_')

                if not(package and release):
                    if debug:
                        print "Didn't find debian package information: ", package, release, deb_pkg
                    continue

                deb_pkg_code += """
  if(isdpkgvuln(pkg:"%s", ver:"%s", rls:"%s"))
  {
    security_hole(0);
    exit(0);
  }\n""" % (package, version, platform)

            return(deb_pkg_code)
        except Exception, msg:
            print "Exception in : GenerateCode -> generateDebPkgCheck() method"
            sys.exit(msg)

    def generateHPUXPkgCheck(self, parse, platform, debug = 0):
        """
        Generates HP-UX PKG verifying code
        """

        try:
            pkg_code = ''

            if debug:
                print "Generating RPM Code..."

            for pkg_info in parse.Packages[platform]:
                patch = []
                revision = ''

                (pkg, apply) = pkg_info.split('~')
                pkg = pkg.strip()

                if not(pkg or apply):
                    if debug:
                        print "Didn't find package or revision/patch information..."
                    continue

                if apply.startswith('PH'):
                    for i in apply.split(','):
                        i = i.strip()
                        patch.append(i) 
                else:
                    revision = apply.strip()

                if revision:
                    pkg_code += """
  if(ishpuxpkgvuln(pkg:"%s", revision:"%s", rls:"%s"))
  {
    security_hole(0);
    exit(0);
  }\n""" % (pkg, revision, platform)

                elif patch:
                    pkg_code += """
  if(ishpuxpkgvuln(pkg:"%s", patch_list:%s, rls:"%s"))
  {
    security_hole(0);
    exit(0);
  }\n""" % (pkg, patch, platform)
                else:
                    print "ERROR: Din't find revision and patch..." 

            return(pkg_code)

        except Exception, msg:
            print "Exception in : GenerateCode -> generateHPUXPkgCheck() method"
            sys.exit(msg)


    def generateGentooPkgCheck(self, parse, platform, debug = 0):
        """
        Generates Gentoo PKG verifying code
        """

        try:
            gentoo_pkg_code = ''
            packages = {}

            if debug:
                print "Generating Gentoo Pkg Code..."

            ## Example for parse.Packages "{'All supported architectures': 
            ## {'kde-base/kdelibs': ['lt', '4.0', 'rge', '3.5.8-r4', 'rge',
            ## '3.5.9-r3', 'gt', '4.0', 'lt', '3.5.5', 'rge', '3.5.10-r2']}}"

            if parse.Packages.has_key(platform):
                packages = parse.Packages[platform]
            else:
                if debug:
                    print "ERROR: %s Platform not found in %s" %(platform, parse)

            for pkgs in packages:
                    ver_list = packages[pkgs]
                    ## ver_list should have even number for elements,
                    ## i.e Key Pair "le","1.2"
                    if len(ver_list)%2 == 1 and len(ver_list) != 3:
                        if debug:
                            print "ERROR: Something wrong in the dict : ", ver_list
                        return {}

                    package = pkgs.strip()
                    vuln_ver_list = []
                    un_aff_list = []

                    ## Create list for vulnerable and unaffected versions
                    flag = 0
                    count = 0
                    for i in ver_list:
                        if "break" in i:
                            count = count + 1
                            flag = 1
                            continue
                        if not re.findall('[0-9.]+', i):
                            if not flag:
                                vuln_ver_list.append(str(i + " " + \
                                                       ver_list[count + 1]))
                            else:
                                un_aff_list.append(str(i + " " + \
                                                       ver_list[count + 1]))

                        count = count + 1

                    if not un_aff_list:
                        un_aff_list = [""]

                    ## unaffected list may be empty, but not package and
                    ## vulnerable version
                    if not(package and vuln_ver_list):
                        if debug:
                            print "ERROR: Didn't find all gentoo package "+ \
                                      "information: ", package, vuln_ver_list
                        continue

                    gentoo_pkg_code += """

if(ispkgvuln(pkg:"%s", vulnerable:%s, unaffected:%s))
{
  security_hole(0);
  exit(0);
}""" % (package, vuln_ver_list, un_aff_list)

            return(gentoo_pkg_code)
        except Exception, msg:
            print "Exception in : GenerateCode -> generateGentooPkgCheck() method"
            sys.exit(msg)


    def generateSolarisPkgCheck(self, parse, platform, debug = 0):
        """
        Generates Solaris PKG verifying code
        """

        try:
            gentoo_pkg_code = ''
            os = release = arch = patch = ''
            packages = {}

            if debug:
                print "Generating Solaris Pkg Code..."

            ## Example for parse.Packages {'solaris_5.10_sparc': 
            ## ['SUNWcnetr', 'SUNWcsl', 'SUNWckr', 'SUNWcsr', 'SUNWhea'],
            ## 'solaris_5.9_sparc': ['SUNWcnetr', 'SUNWcsl', 'SUNWckr', 
            ## 'SUNWcsr', 'SUNWhea']} 

            if not platform:
                if debug:
                    print "ERROR: Platform not Found", platform
                return ''

            all_fields = platform.split('_')
            if len(all_fields) == 3:
                release = all_fields[1]
                if all_fields[2] == 'x86':
                    arch = 'i386'
                elif all_fields[2] == 'sparc':
                    arch = 'sparc'
                else:
                    if debug:
                        print "ERROR: x86/sparch Architecture not found : ", \
                                                                    all_fields
                    return ''
            else:
                if debug:
                    print "ERROR: Required Fields not found os=%s, release=%s," \
                                                   " arch=%s" %(os, release, arch)
                return ''

            if parse.Packages.has_key(platform):
                if len(parse.Packages[platform]) == 2:
                    packages, patch = parse.Packages[platform]
                else:
                    if debug:
                        print "ERROR: Packages or Patch not found in ", \
                                                 parse.Packages[platform]
                    return ''
            else:
                if debug:
                    print "ERROR: %s Platform not found in %s" \
                                      %(platform, parse.Packages)
                return ''

            gentoo_pkg_code += """

if(solaris_check_patch(release:"%s", arch:"%s", patch:"%s", package:"%s") < 0)
{
  security_hole(0);
  exit(0);
}""" % (release, arch, patch, packages)

            return(gentoo_pkg_code)
        except Exception, msg:
            print "Exception in : GenerateCode -> generateSolarisPkgCheck() "+ \
                                                                   "method", msg
            sys.exit(msg)


class GenerateNVTLocal:
    """
    GenerateNVTLocal: NASL code generator based on the template
    """

    ActionObj = GenerateCode()

    func_dict = {

        '___IF_RELEASE_OPEN___'   : 'self.ActionObj.generateReleaseCheck\
                                                 (platform, debug)',

        '___IS_RPM_VULN___'       : 'self.ActionObj.generateRPMCheck\
                                     (self.parse, platform, debug)',

        '___IS_DEB_PKG_VULN___'       : 'self.ActionObj.generateDebPkgCheck\
                                     (self.parse, platform, debug)',

        '___IS_HPUX_PKG_VULN___'       : 'self.ActionObj.generateHPUXPkgCheck\
                                     (self.parse, platform, debug)',

        '___IS_GENTOO_PKG_VULN___'       : 'self.ActionObj.generateGentooPkgCheck\
                                     (self.parse, platform, debug)',

        '___IS_SOLARIS_PKG_VULN___'       : 'self.ActionObj.generateSolarisPkgCheck\
                                     (self.parse, platform, debug)',

        '___IF_RELEASE_CLOSE___'  : 'self.ActionObj.generateClose(debug)',
    }


    def _truncateDescription(self, debug=0):
        """
        Truncate the description, if it's too long and append message
        """
        full_desc = self.parse.Description + self.parse.Impact + \
                         self.parse.Platforms + self.parse.Product
        if len(full_desc) > 2848:
            trun_len = len(self.parse.Description) - (len(full_desc) - 2660)
            description = self.parse.Description[0:trun_len]
            description += " ... \n\n  Description truncated, for more " +\
                           "information please check the Reference URL"
            if debug:
                print "Description Truncated..."
        else:
            description = self.parse.Description

        return description


    def addScriptTags(self, template, cve_list, debug):
        """
        Get cvss score by visiting NVD and calculate risk factor based on cvss base score
        """
        ## Default Risk Factor is high 
        risk_fact = "High"
        highestCVSSScore = 0

        if cve_list:
            ## Convert to list
            cve_list = self.parse.CVEs.split('", "')

            ## Get Highest cvss base score from list of cve's
            highestCVSSScore = utils.getHighestCVSSScoreFromCVEList(cve_list, debug)

            ## Calculate risk factor based on cvss base score
            ## default is risk factor is "High". i.e if cvss score not present
            if highestCVSSScore:
                risk_fact = utils.calculateRiskFactor(highestCVSSScore, debug)
                template = string.replace(template, '__CVSS_BASE_SCORE__', \
                                                         highestCVSSScore)
        if not highestCVSSScore:
            template = re.sub('\s?\s?script_tag\(name:"cvss_base", value:"__CVSS_BASE_SCORE__"\);\n', '', template)

        template = string.replace(template, '__RISK_FACTOR__', risk_fact)

        return template


    def generateCode(self, template, script_id , reference, debug=0):
        """
        Reads the template and according to the code generator constants in the
        template, the corresponding generator function is called to create the NVT
        code.
        template - code skeleton
        script_id - New script ID
        reference - Reference URL of the advisory, for which the code is being
        generated
        """

        try:
            main_code = ''
            final_code = ''
            if_open_code = ''
            if_close_code = ''
            final_pkg_code = ''

           ## If parsed data doesn't contain affected packages
            if not self.parse.Packages:
                return None

            genConst = re.findall("___.*___", template)
            if debug:
                print "Special Constants from Template : ", genConst

            ## Creates the code portion of NVT
            for platform in self.parse.Packages.keys():
                flag = 0
                for i in genConst:
                    if "OPEN" in i and self.func_dict.has_key(i):
                        if_open_code = eval(self.func_dict[i])
                        flag = 1
                    elif "CLOSE" in i and self.func_dict.has_key(i):
                        if_close_code = eval(self.func_dict[i])
                    elif self.func_dict.has_key(i):
                        main_code += eval(self.func_dict[i])

                    if (if_open_code and if_close_code and main_code):
                        final_pkg_code = final_pkg_code + '\n' + if_open_code \
                                                    + main_code + if_close_code
                        main_code = if_open_code = if_close_code = ''
                        flag = 0
                    elif (main_code and (not flag)):
                        final_pkg_code = final_pkg_code + main_code
                        main_code = ''

                final_code = final_code + final_pkg_code
                final_pkg_code = ''

            if debug:
                print "Code generated for Script ID : ", script_id

            ## Remove unwanted stuff from template
            for i in genConst:
                template = string.replace(template, i, "")

            ## Truncating description, if it's too long.
            description = self._truncateDescription(debug)

            ## Replace all "__" Veriables in the template with the parsed content

            template = string.replace(template, '__SCRIPT_REF__', reference)

            template = string.replace(template, '__SCRIPT_ID__', script_id)

            template = string.replace(template, '__SCRIPT_DESC__', description)

            ## Remove CVEIDS line, If CVEs are not present.
            if not self.parse.CVEs:
                template = re.sub('\s?\s?script_cve_id\("__CVEIDS__"\);.*?\n','', template)
            else:
                template = string.replace(template, '__CVEIDS__', \
                                                           self.parse.CVEs)
   
            ## Replace cvss base score and risk factor 
            template = self.addScriptTags(template, self.parse.CVEs, debug)

            template = string.replace(template, '__SCRIPT_IMPACT__', \
                                                            self.parse.Impact)

            template = string.replace(template, '__SCRIPT_PLAT__', \
                                                         self.parse.Platforms)

            template = string.replace(template, '__SCRIPT_PROD__', \
                                                           self.parse.Product)

            template = string.replace(template, '__SCRIPT_NAME__', \
                                                              self.parse.Name)

            template = string.replace(template, '__XREF_NAME__', \
                                                          self.parse.XREF[0])

            template = string.replace(template, '__XREF_VALUE__', \
                                                          self.parse.XREF[1])

            template = string.replace(template, '__SCRIPT_PKG__', \
                                                           self.parse.Summary)

            ## Remove unwanted lines and spaces
            template = template.strip()

            ## Append main code protion of the script
            template = template + final_code

            ## Remove unwanted lines and spaces
            template = template.strip()

            return(template)

        except Exception, msg:
            print "Exception in : createscript -> GenerateNVTLocal class"
            sys.exit(msg)
