###############################################################################
# Local Security Check Automation Framework
#
# Authors:
# Veerendra GG <veerendragg@secpod.com>
#
# Revision 1.0
# Date: 2008/12/24
#
# 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.
###############################################################################

import sys
import utils


## Read, write pickle object file and Get, Set the Script ID
class IDRecorder:
    """
    IDRecorder: Record the Script ID's used in NVT's

    Functions:
    - getID(): Obtain New ID
    - recordID(): Cache the used ID
    - recordFinal(): Dump the IDs to cache
    """

    id_file = ''
    start_id = ''
    pickle_id_dict = {}


    def __init__(self, debug=0):
        """
        Loads the pickle object, If not proper builds new one
        """

        try:
            ## Read and validate pickle object file
            self.pickle_id_dict = utils.loadPickle(self.id_file, debug)

            if not self.pickle_id_dict:
                self.pickle_id_dict = {}
                ## Regenerate proper pickle object file, if it deleted or altered
                if debug:
                    print "\t IDFile and scripts are getting regenerated ...\n"

                file_handle = open(self.id_file, 'w+')
                self.pickle_id_dict['startID'] = int(self.start_id) - 1
                file_handle.close()

            if(debug):
                print 'ID Dictonary Contains : ', self.pickle_id_dict

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


    def recordFinal(self, debug=0):
        """
        Write pickle object to file
        """

        if debug:
            print "Caching the Script ID's to file..."
        utils.dumpPickle(self.pickle_id_dict, self.id_file, debug)


    def recordID(self, new_id, link_id, debug=0):
        """
        Record script id, given the newid and linkid
        newid - New script ID
        link_id - a unique string quanifying the advisory
        """

        try:
            ## The script is already generated for this link_id
            if self.pickle_id_dict.has_key(link_id):
                return None

            if debug:
                print "Recording the New Id ..."
                print "Link ID : (%s) : New ID : (%s)"%(link_id, new_id)
            self.pickle_id_dict[link_id] = new_id

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


    def getID(self, link_id, debug=0):
        """
        Obtain new script ID for a given link_id
        link_id - a unique string quantifying the advisory
        """

        try:
            ## The script is already generated for this link_id
            if self.pickle_id_dict.has_key(link_id):
                return None

            if debug:
                print "New Id is getting generated ..."

            idList = self.pickle_id_dict.values()
            idList.sort()
            if int(self.start_id) in idList:
                new_id = int(idList[-1]) + 1
            else:
                new_id = int(self.start_id)

            return new_id

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