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


def readFile(filename, debug=0, lines=0):
    """
    Read the given file, if lines=1 returns list
    """

    try:
        if debug:
            print "Reading from file :", filename
        file_handle = open(filename, 'r')

        if lines:
            data = file_handle.readlines()
        else:
            data = file_handle.read()

        file_handle.close()
        return data

    except Exception, msg:
        print "Error in Reading File : ", filename
        print "Exception in : utils -> readFile() method."
        sys.exit(msg)


def writeFile(filename, data,  debug=0):
    """
    Write data to the given file
    """

    try:
        if debug:
            print "Writing to file :", filename
        file_handle = open(filename, 'w')
        file_handle.write(data)
        file_handle.close()

    except Exception, msg:
        print "Error in Writing to file : ", filename
        print "Exception in : utils -> writeFile() method."
        sys.exit(msg)


def fetchFiles(url, file_name, debug=0):
    """
    Retrive the content from the given URL, and save to the specified cache
    directory
    """

    try:
        url_open = urllib.URLopener()

        if debug:
            print "Retrieving content from : ", url

        url_open.retrieve(url, file_name)

        return file_name

    except Exception, msg:
        sys.exit(msg)


def stringFormat(file, debug=0):
    """
    Returns properly formatted string
    """

    lines = ''

    try:
        if debug:
            print "Reading and Formatting File : ", file

        lines = readFile(file, debug, lines=1)
        lines = "".join(lines)
        return(lines)

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


def removeDups(list_):
    """
    Remove duplicates from a list
    """

    tmp = []
    if not list_:
        return tmp

    for i in list_:
       i = i.strip()
       if i not in tmp:
           tmp.append(i)

    return tmp


def stripIt(strip_list, strip_val):
    """
    Returns Stripped list from list
    """
    if not strip_list:
        return []

    striped_list = []
    if type(strip_val) == list:
        for i in strip_list:
            for j in strip_val:
                i = i.replace(j, '')
            striped_list.append(i)
    elif type(strip_val) == str:
        for i in strip_list:
            striped_list.append(i.replace(strip_val, ''))

    return striped_list


def getHTMLCon(url, debug=0):
    """
    Returns Contents, Given URL
    """

    if not url:
        return None
    ptr = urllib.urlopen(url)
    data = ptr.read()
    ptr.close()
    return data


def loadPickle(file_name, debug=0):
    """
    Load pickle file, given path
    """
    if not file_name:
        return None
    try:
        return pickle.load(open(file_name, 'r'))
    except Exception, msg:
        return None


def dumpPickle(data, file_name, debug=0):
    """
    Dump object into the file
    """
    try:
        pickle.dump(data, open(file_name, 'w'))
        return True
    except Exception, msg:
        return False


def formatMultiLines(data, debug=0):
    """
    Format the lines to 79 characters
    """

    tmp = ''
    formated_data = ''

    data = data.strip()
    for i in data.split():
        if (len(tmp) + len(i)) >= 77:
          formated_data += tmp + '\n' + '  '
          tmp = ''
        tmp = tmp + i + ' '

    formated_data += tmp

    return formated_data


def createDir(path, debug=0):
    """
    Create Directory, If not exists
    """

    if not path:
        if debug:
            print 'Path not passed to _createDir() Function'
        return None

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

    return True


def getCVSSFromNVDWebSite(cve, debug=0):
    """
    Get CVSS Score from NVD Site
    """
    try:
        import re
        url = "http://web.nvd.nist.gov/view/vuln/detail?vulnId="
        url = url + cve

        ## Get content from given url
        nvd_data = getHTMLCon(url, debug)

        if not nvd_data:
            return None

        if "CVE was not found" in nvd_data:
            if debug:
                print "CVE not found : ", url
            return None

        ## Get CVSS Score from NVD
        base_score = re.findall("CVSS v2 Base Score:.*>(\d\d?.\d)<",nvd_data)
        if not base_score:
            if debug:
                print "CVSS Score not present : ", url
            return None

        return base_score[0]

    except Exception, msg:
        print "Error in while getting CVSS score from NVD from \n : ", msg
        print "Exception in : getCVSSFromNVDWebSite() method."
        sys.exit(msg)


def getHighestCVSSScoreFromCVEList(cve_list, debug=0):
    """
    Return Highest Base Score from CVE List by visiting each CVE
    """
    try:
        if not cve_list:
            if debug:
                print "CVE List is Empty"
            return None

        highest_cvss_score = 0

        if debug:
            print "\nGetting CVSS Score from NVD for CVE's : ", cve_list

        ## Visit each CVE and get highest CVSS Score
        for cve in cve_list:
            ## Get CVSS score from CVE
            cvss_score = getCVSSFromNVDWebSite(cve, debug)
            if not cvss_score:
                print "Ignoring CVE, as CVSS Score not present : ", cve
                continue

            if debug:
                print "CVSS Score for %s is %s "%(cve, cvss_score)

            ## Calculate highest CVSS Score
            if float(cvss_score) > float(highest_cvss_score):
                highest_cvss_score = cvss_score

        if highest_cvss_score == 0:
            print "ERROR: CVSS Score nor present for CVE List : ", cve_list
            return None

        return highest_cvss_score

    except Exception, msg:
        print "Error in while getting highest CVSS score from NVD from \n : ", msg
        print "Exception in : getHighestCVSSScoreFromCVEList() method."
        sys.exit(msg)


def calculateRiskFactor(cvss_score, debug=0):
    """
    Calculates and Return Risk Factor given CVSS Base Score/Temporal Score
    """
    try:
        risk_factor = ""

        if not cvss_score:
            if debug:
                print "ERROR: CVSS Score is Null"
            return None

        if debug:
            print "Calculating Risk Factor based on given CVSS Base Score : ", cvss_score
        cvss_score = float(cvss_score)

        ## Calculate Risk Factor
        if cvss_score >= 0 and cvss_score <= 2:
            risk_factor = "Low"
        elif cvss_score > 2 and cvss_score <= 5:
            risk_factor = "Medium"
        elif cvss_score > 5 and cvss_score <= 8:
            risk_factor = "High"
        elif cvss_score > 8 and cvss_score <= 10:
            risk_factor = "Critical"

        return risk_factor

    except Exception, msg:
        print "Exception while Calculating Risk Factor : ", msg
        print "Exception in : calculateRiskFactor() method."
        sys.exit(msg)
