vulnerability-reportv9.py

import requests
import datetime
from openpyxl import Workbook, load_workbook
from openpyxl.styles import PatternFill
import json
import os
import logging

# Authenticate to Prisma Cloud and get the token
def authenticate():
    url = "https://api.prismacloud.io/login"
    payload = json.dumps({
        "customerName": "Tenant Name",
        "username": "Access Key",
        "prismaId": "Prisma ID",
        "password": "Secret Key"
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.post(url, headers=headers, data=payload)

    if response.status_code == 200:
        print("Authentication successful.")
        token = response.json().get('token')
        if not token:
            raise Exception("No token received.")
        return token
    else:
        raise Exception(f"Failed to authenticate: {response.status_code}, {response.text}")

# Fetch data from the Vulnerability API using the generated auth token
def fetch_data(api_url, token):
    headers = {
        'Authorization': f"Bearer {token}",
        'accept': 'application/json',
        'content-type': 'application/json',
        'x-prisma-route': '/dashboards/vulnerabilities',
        'x-redlock-auth': token
    }

    response = requests.get(api_url, headers=headers)
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"Failed to fetch data: {response.status_code}, {response.text}")

# Create a fill pattern for coloring cells
def get_fill(color):
    return PatternFill(start_color=color, end_color=color, fill_type="solid")

# Creates/Updates the workbook with daily data and summary
def update_workbook(data, filename):
    try:
        workbook = load_workbook(filename)
    except FileNotFoundError:
        workbook = Workbook()
        workbook.remove(workbook.active)  # Remove default sheet if creating new workbook

    overview = data['overviewSummary']
    total_vulnerabilities = overview['totalVulnerabilitiesinRuntime']
    total_remediated = overview['totalRemediatedinRuntime']

    # Prepare today's sheet
    today = datetime.date.today()
    sheet_name = today.strftime("%Y-%m-%d")

    if sheet_name not in workbook.sheetnames:
        daily_sheet = workbook.create_sheet(sheet_name)
        daily_sheet.append(["Total Count", "Critical Count", "High Count", "Medium Count", "Low Count", "Critical Remediation Count", "High Remediation Count", "Medium Remediation Count", "Low Remediation Count"])
        daily_sheet.append([
            total_vulnerabilities["totalCount"],
            total_vulnerabilities["criticalCount"],
            total_vulnerabilities["highCount"],
            total_vulnerabilities["mediumCount"],
            total_vulnerabilities["lowCount"],
            total_remediated["criticalCount"],
            total_remediated["highCount"],
            total_remediated["mediumCount"],
            total_remediated["lowCount"]
        ])

        # Apply colors to the cells
        for cell in daily_sheet["2:2"]:
            if cell.column == 2 or cell.column == 6:
                cell.fill = get_fill("FF0000")  # Red for critical
            elif cell.column == 3 or cell.column == 7:
                cell.fill = get_fill("FFA500")  # Orange for high
            elif cell.column == 4 or cell.column == 8:
                cell.fill = get_fill("FFFF00")  # Yellow for medium
            elif cell.column == 5 or cell.column == 9:
                cell.fill = get_fill("00FF00")  # Green for low

    # Update summary sheet
    if "Summary" not in workbook.sheetnames:
        summary_sheet = workbook.create_sheet("Summary")
        summary_sheet.append(["Date", "Total Count", "Critical Count", "High Count", "Medium Count", "Low Count", "Critical Remediation Count", "High Remediation Count", "Medium Remediation Count", "Low Remediation Count"])

    summary_sheet = workbook["Summary"]
    summary_sheet.append([
        sheet_name,
        total_vulnerabilities["totalCount"],
        total_vulnerabilities["criticalCount"],
        total_vulnerabilities["highCount"],
        total_vulnerabilities["mediumCount"],
        total_vulnerabilities["lowCount"],
        total_remediated["criticalCount"],
        total_remediated["highCount"],
        total_remediated["mediumCount"],
        total_remediated["lowCount"]
    ])

    # Apply colors to the summary sheet cells
    for cell in summary_sheet[f"{summary_sheet.max_row}:{summary_sheet.max_row}"]:
        if cell.column == 3 or cell.column == 7:
            cell.fill = get_fill("FF0000")  # Red for critical
        elif cell.column == 4 or cell.column == 8:
            cell.fill = get_fill("FFA500")  # Orange for high
        elif cell.column == 5 or cell.column == 9:
            cell.fill = get_fill("FFFF00")  # Yellow for medium
        elif cell.column == 6 or cell.column == 10:
            cell.fill = get_fill("00FF00")  # Green for low

    workbook.save(filename)

# Create a log file to save the daily execution
def create_log_file(success, error_message=None):
    log_directory = "./logs"
    os.makedirs(log_directory, exist_ok=True)
    today = datetime.date.today()
    log_file_name = f"vulnerability-report-log-{today}.txt"
    log_file_path = os.path.join(log_directory, log_file_name)
    with open(log_file_path, "w") as log_file:
        if success:
            log_file.write("Script executed successfully.\n")
        else:
            log_file.write(f"Script failed with error: {error_message}\n")

# Main function that runs the daily process
def main():
    try:
        auth_token = authenticate()
        print(f"Auth Token: {auth_token}")
        api_url = "https://api.prismacloud.io/uve/api/v2/dashboard/vulnerabilities/overview"
        today = datetime.date.today()

        # Generate filename based on the current month
        filename = f"vulnerabilities_report_{today.strftime('%Y-%m')}.xlsx"

        # Check if the month has changed since the last execution
        last_execution_month = None
        try:
            with open("last_execution_month.txt", "r") as file:
                last_execution_month = file.read().strip()
        except FileNotFoundError:
            pass

        current_month = today.strftime("%Y-%m")
        if last_execution_month != current_month:
            # Create a new spreadsheet file for the new month
            filename = f"vulnerabilities_report_{current_month}.xlsx"
            with open("last_execution_month.txt", "w") as file:
                file.write(current_month)

        data = fetch_data(api_url, auth_token)
        update_workbook(data, filename)
        create_log_file(success=True)
    except Exception as e:
        create_log_file(success=False, error_message=str(e))
        print(f"An error occurred: {e}")

if __name__ == "__main__":
    main()