#!/usr/bin/python
###############################################################################
# Local Security Check Automation Framework
#
# Authors:
# Veerendra GG <veerendragg@secpod.com>
# Date Written: 2008/12/15
#
# Revision 1.1
#
# Copyright:
# Copyright (c) 2009 SecPod , http://www.secpod.com
#
# 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.
###############################################################################

"""
LSCGenerator: Local Security Check Generator
Main module that reads configuration values and sets to the respective modules.
Dependening on the value set in 'generate' of lsc.conf, the parsers and
the generator is invoked to produce the NVT's.
"""


import os
import sys
import commands
import ConfigParser

from common import *
from parser import *
from test import sanity_test

## Set debug to 1, to run in debug mode
debug = 0

## set the config filename
config_file = 'lsc.conf'

class LSCGenerator:

    ## Map between 'generate' conf value and the corresponding parser object
    ## and configuration section.
    ## Each new parser added, will have to be updated here. Better to move this
    ## to parser package?
    _object_map = {
        'Suse'    : (suse, 'SUSE'),
        'Fedora'  : (fedora, 'FC'),
        'Redhat'  : (redhat, 'Redhat'),
        'Centos'  : (centos, 'CentOS'),
        'Ubuntu'  : (ubuntu, 'Ubuntu'),
        'Mandriva'  : (mandriva, 'Mandriva'),
        'Hpux'  : (hp_ux, 'HPUX'),
        'Solaris'  : (solaris, 'Solaris'),
        'Debian'  : (debian, 'Debian'),
        'Gentoo'  : (gentoo, 'Gentoo'),

    }


    def __init__(self, debug=0):
        """ Initialize configuration reader """

        self.cwd = os.getcwd()
        self.mantis_obj = None
        self.config = ConfigParser.ConfigParser()
        if(debug):
            print "Current Working Dir : ", self.cwd


    def setUp(self, path):
        """ Create initial required setup"""

        self._createDir(path)
        self.idfile_path = path + 'IdCache/'
        self.html_cache_path = path + 'HtmlCache/'
        self.nasl_script_path = path

        self._createDir(self.idfile_path)
        self._createDir(self.html_cache_path)
        for i in self._object_map.keys():
            i = i.title().strip()
            self._createDir(self.html_cache_path + i)
            self._createDir(self.nasl_script_path + i)


    def _createDir(self, path):
        """ Create Directory, If not exists """
        if not path:
            print 'Path not passed to _createDir() Function'
            return None

        if not os.path.exists(path):
            os.mkdir(path)


    def _readConfAndGetBuildPath(self, debug=0):
        """ Read the build path from config file """
        try:
            ## Reading from Configuration file
            ConfigFile = os.path.join(self.cwd, config_file)
            if debug:
                print 'Config file path : ', ConfigFile

            config_open = open(ConfigFile, 'r+')
            self.config.readfp(config_open)

            ## Read and Set build path
            build_path = ''
            build_path = self.config.get('common', 'build_path').strip()
            if not build_path:
               build_path = self.cwd

            if not build_path.endswith('/'):
                build_path = build_path + '/'

            return build_path

        except Exception , msg:
            print "Exception in : LSCGenerator -> "+ \
                                           "_readConfAndGetBuildPath() method"
            print "Error in reading build path...", msg
            sys.exit(msg)


    def _sanityTestSetUp(self, sanity_test_obj, debug=0):
        """ Read and set required attributes for Sanity test"""
        try:

            build_path = self._readConfAndGetBuildPath(debug)
            ## Create Required Dir
            self._createDir(build_path + 'sanity_test')
            self._createDir(build_path + 'sanity_test/plugins/')

            ## Read Sanity test config paths
            openvas_plugin_path = self.config.get('test', \
                                            'openvas_plugin_path').strip()
            openvas_bin_path = self.config.get('test', \
                                               'openvas_bin_path').strip()

            if not (openvas_plugin_path or openvas_bin_path):
                print "test section Config Values are missing."

            os_list = self._object_map.keys()

            ## Set Attributes
            setattr(sanity_test_obj, 'cwd', self.cwd)
            setattr(sanity_test_obj, 'build_path', build_path)
            setattr(sanity_test_obj, 'os_list', os_list)
            setattr(sanity_test_obj, 'openvas_plugin_path',openvas_plugin_path)
            setattr(sanity_test_obj, 'openvas_bin_path', openvas_bin_path)

        except Exception, msg:
            print "Exception in : LSCGenerator -> _sanityTestSetUp() method"
            rm_cmd = 'rm -rf ' + build_path + 'sanity_test'
            commands.getoutput(rm_cmd)
            sys.exit(msg)


    def _performSanityTest(self, sanity_test_obj, debug=0):
        """ Perform Sanity Test, by compiling the scripts """
        try:
            build_path = self._readConfAndGetBuildPath(debug)
            if debug:
                print "####################################################################"
                print "Performing sanity test..."

            ## check for script errors, by compileing them
            msg = sanity_test_obj.compileScripts(debug)
            if msg:
                print msg

            if debug:
                print "Sanity test complete."
                print "####################################################################"

            ## Remove command
            rm_cmd = 'rm -rf ' + build_path + 'sanity_test'

            ##Clean Up
            commands.getoutput(rm_cmd)

        except Exception, msg:
            print "Exception in : LSCGenerator -> _performSanityTest() method"
            rm_cmd = 'rm -rf ' + build_path + 'sanity_test'
            commands.getoutput(rm_cmd)
            sys.exit(msg)


    def _mantisReportSetUp(self, debug=0):
        """ Read and Set Attributes for Creating Mantis Report. """
        try:
            if debug:
                print "Getting mantis config vaule..."

            conf_sec = 'mantis'
            host = self.config.get(conf_sec, 'db_host').strip()
            db_name = self.config.get(conf_sec, 'dbname').strip()
            mysql_passwd = self.config.get(conf_sec, 'mysql_passwd').strip()
            project_name = self.config.get(conf_sec, 'project_name').strip()
            mysql_user_name = self.config.get(conf_sec, 'mysql_user').strip()
            mantis_user_name = self.config.get(conf_sec, 'mantis_user').strip()

            if debug:
                print "Setting mantis attributes for mantis task creation..."

            setattr(self.mantis_obj, 'host', host)
            setattr(self.mantis_obj, 'db_name', db_name)
            setattr(self.mantis_obj, 'mysql_passwd', mysql_passwd)
            setattr(self.mantis_obj, 'project_name', project_name)
            setattr(self.mantis_obj, 'mysql_user_name', mysql_user_name)
            setattr(self.mantis_obj, 'mantis_user_name', mantis_user_name)

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


    def execute(self, debug=0):
        """ Based on the 'generate' config value, generate() is invoked
            with the right parser object.
        """
        try:
            ## Get Build Path
            build_path = self._readConfAndGetBuildPath(debug)

            if debug:
                print "Everyting will be generated in the Path : ", build_path

            ## Create Initial Setup
            if debug:
                print "#####################################################################"
                print "LSC Generator configuration..."
            self.setUp(build_path)

            ## Get the list of Operating Systems for which LSC have to
            ## be developed
            genList = self.config.get('common', 'generate')
            if not genList:
                print "No value given for 'generate' in config : ", \
                                                          config_file
                return 0
            genList = genList.split(',')

            ## Create Mantis task, If enabled
            mantis_enabled = self.config.get('mantis', \
                                             'generate_mantis_report').strip()
            mantis_enabled = mantis_enabled.strip().title()

        except Exception, msg:
            print 'Config ERROR: (%s), Check %s Conf file' %(msg, config_file)
            sys.exit(msg)

        try:
            ## Create mantis task, If mantis_enabled is Yes
            if mantis_enabled == 'Yes':
                ## Create Object of MantisReporter
                from report import mantisreporter
                self.mantis_obj = mantisreporter.MantisReporter()

                ## Read conf and set the variables for future use.
                self._mantisReportSetUp(debug)

            ## Invoke the generate(), passing the parser object
            for parser_ref in genList:
                if not parser_ref:
                    continue
                parser_ref = parser_ref.title().strip()

                if self._object_map.has_key(parser_ref):
                    ## Set proper Paths
                    paths = {}
                    paths['id_cache_path'] = self.idfile_path  + parser_ref
                    paths['html_cache_path'] = self.html_cache_path+parser_ref \
                                                                    + '/'
                    paths['nasl_script_path'] = self.nasl_script_path \
                                                      + parser_ref + '/'

                    ## Get Parser Object and Config Value
                    (parser_obj, config_sect) = self._object_map[parser_ref]

                    print "#################################################################"
                    print "Generating the code for %s ... "%(config_sect)
                    self.generate(parser_obj, config_sect, paths, debug)
                    print "Successfully generated the Local Security Checks"
                    print "#################################################################"
                else:
                    print "The parser for %s doesn't exist" %(parser_ref)

            print "Local Security Checks are generated."

            ## Sanity Testing
            if debug:
                print "#################################################################"
                print "Preparing to perform sanity test..."
            sanity = self.config.get('test', 'sanity_test').strip()
            if not sanity:
                sanity = 'No'

            if sanity.title() == 'Yes':
                sanity_test_obj = sanity_test.SanityTest()
                self._sanityTestSetUp(sanity_test_obj, debug)
                self._performSanityTest(sanity_test_obj, debug)

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


    def generate(self, parser_obj, config_sect, paths, debug=0):
        """ Core function, does,
            - read and set config values
            - invoke parser to fetch and parse HTML's
            - invoke script generator
            - manage Script ID's
        """

        try:
            ## Set Paths
            id_file = paths['id_cache_path']
            html_cache = paths['html_cache_path']
            script_directory = paths['nasl_script_path']

            ## Read from conf
            year = self.config.get(config_sect, 'year').strip()
            main_url = self.config.get(config_sect, 'mainurl').strip()
            template = self.config.get(config_sect, 'template').strip()
            start_id = self.config.get(config_sect, 'startid').strip()

            if not (year and html_cache and main_url and template and \
                            script_directory and id_file and start_id):
                print "Mandatory config value is missing"
                sys.exit(0)

            if debug:
                print "Configuration complete."
                print "#####################################################################"

        except Exception, msg:
            print 'Unable to get mandatory values from %s Config file'\
                                                         %(config_file)
            sys.exit(msg)

        try:
            parse = parser_obj.Parser()
            generator = generate_script.GenerateNVTLocal()

            ## Set the Required attributes
            setattr(generator, 'parse', parse)
            setattr(parse, 'html_cache', html_cache)
            setattr(parse, 'main_url', main_url)
            setattr(record_id.IDRecorder, 'start_id', start_id)
            setattr(record_id.IDRecorder, 'id_file', id_file)

            RecordID = record_id.IDRecorder(debug)
            read_template = utils.readFile(template, debug)

            ## Retrieve HTML, if not present in cache
            if debug:
                print "#######################################################################"
                print "Fetching HTML advisories"

            parse.fetchHTML(year, debug)

            if debug:
                print "Advisory fetching complete"
                print "#######################################################################"

            for cache_file in os.listdir(html_cache):
                if debug:
                    print "#####################################################################"
                    print "Parsing the cached file and generating the check"+\
                                                          " for " + cache_file
                link_id = cache_file.split('.')[0]

                ## Get New script ID
                new_id = RecordID.getID(link_id, debug)
                if not new_id:
                    if debug:
                        print 'Script ID not available, script is already ' +\
                                                 'generated for : ' + link_id
                    continue
                if debug:
                    print "New Script ID is : ", new_id

                ## Call Parser and Generate the code
                html_content = utils.stringFormat(html_cache + cache_file, \
                                                               debug)
                parse_success = parse.parser(html_content, debug)
                if not parse_success:
                    continue;

                ## Generate the final code
                reference = parse.formatReference(main_url, cache_file)
                final_template = generator.generateCode(read_template, \
                                 str(new_id), reference, debug)
                ## If Packages(self.parse.Packages) is empty
                if not final_template:
                    continue

                ## Convert to ascii/iso-8859-1 encoding.
                try:
                    final_template = final_template.encode('ascii')
                except UnicodeDecodeError:
                    try:
                        final_template = final_template.decode('utf-8')
                        final_template = final_template.encode('iso-8859-1')
                    except UnicodeDecodeError:
                        print "Warning: Found encoding other then ascii/utf-8 .."
                        final_template = final_template.decode('ascii','ignore')
                        final_template = final_template.encode('ascii','ignore')

                ## Script file name
                file_name = 'gb_' + parse.FileName + '.nasl'
                file_name = file_name.replace(' ','_')

                if(debug):
                    print "Generated New ID  %s , Script Name %s " %(new_id, \
                                                                    file_name)
                ## Write generated code to file.
                utils.writeFile(script_directory + file_name, \
                                                   final_template, debug)

                ## Record the New Script ID
                RecordID.recordID(new_id, link_id, debug)

                if debug:
                    print "Generated the check " + file_name
                    print "#############################################################"

                ## Create mantis task
                if self.mantis_obj:
                    ## Create Mantis Task, If it's not created earlier.
                    pic_file_name = self.cwd + "/report/mantis_id.pickle"
                    file_id_dict = utils.loadPickle(pic_file_name, debug)
                    if file_id_dict:
                        if file_id_dict.has_key(file_name):
                            if debug:
                                print "Mantis Task is already created : ", \
                                                                   file_name
                                print "#############################################################"
                            continue
                    else:
                        file_id_dict = {}

                    if debug:
                        print "Creating the Mantis Task " + file_name
                        print "#############################################################"

                    if parse.CVEs:
                        cves = "\n  CVES Are : %s \n" % parse.CVEs
                    else:
                        cves = "\n  CVES Not Present \n"

                    script_id = '  Script ID : %s \n' % new_id
                    script_name = '  Script Name : %s \n' % file_name
                    script_advid = '  Advisories ID : %s \n' % parse.AdvID
                    script_ref = '  Reference : %s \n' % reference

                    aditional_info = script_id + script_name + script_advid + \
                                     cves + script_ref
                    description = parse.Description.strip()
                    summary = parse.Name

                    ## Set the attributes for mantis task creation
                    setattr(self.mantis_obj, 'summary', summary)
                    setattr(self.mantis_obj, 'description', description)
                    setattr(self.mantis_obj, 'aditional_info', aditional_info)

                    ## Create mantis task.
                    mantis_id = self.mantis_obj.execute(debug)
                    if mantis_id:
                        file_id_dict[file_name] = mantis_id
                        utils.dumpPickle(file_id_dict, pic_file_name)
                    else:
                        if debug:
                            print "Mantis Task is not created for : ",file_name

            ## Dump the recorded ID's to serialized cache
            RecordID.recordFinal(debug)

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


def helpMsg(msg=None):
    if msg:
        print 'ERROR : ', msg

    commonMsg()
    print "Before running Local Security Check Framework, configure lsc.conf."
    print "For more information, refer inline comments in lsc.conf"

    print "\nSpecify --help, to get help"
    print "Specify --verbose, to run in debug mode."
    print "Specify --sanity, to perform only sanity test."
    print "Specify --sanity --verbose, to perform only sanity test in debug mode."


def commonMsg():
    print "#############################################################################"
    print "You are Running LSC Generator Framework."
    print "Local Security Check Generator Framework for Security Advisory of different Vendor."
    print "For ex: SUSE Security Advisory, Fedora Security Advisory etc"
    print "#############################################################################"


if __name__ == '__main__':
    try:
        try:
            import getopt
            opts,args = getopt.getopt(sys.argv[1:], '', \
                    ['verbose', 'help', 'sanity'])
        except getopt.error, msg:
            print 'Valid options are : (--help , --verbose and --sanity)'
            sys.exit('Exiting')

        sanity_only = False
        for opt,arg in opts:
            if opt == '--help':
                helpMsg()
                sys.exit('Exiting...')
            if opt == '--verbose':
                debug = 1
                commonMsg()
                print "Running in verbose mode ..."
            if opt == '--sanity':
                sanity_only = True

        lsc_Obj = LSCGenerator(debug)
        ## Perform only sanity test
        if sanity_only:
            print "#################################################################"
            print "Preparing to perform sanity test..."
            sanity_test_obj = sanity_test.SanityTest()
            lsc_Obj._sanityTestSetUp(sanity_test_obj, debug)
            lsc_Obj._performSanityTest(sanity_test_obj, debug)
        else:
            lsc_Obj.execute(debug)

    except Exception, msg:
        print "Message :", msg
