1773811980

Build # 1773811980

Build Type

push

github

Committed by GitHub

Commit Message

Merge pull request #114 from cisagov/dj-WIP-pages

Run Details

18 of 139 branches covered (12.95%)

Branch coverage included in aggregate %.

36 of 254 new or added lines in 6 files covered. (14.17%)

37 existing lines in 3 files now uncovered.

139 of 766 relevant lines covered (18.15%)

0.91 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

29.69

/src/pe_reports/report_generator.py

"""cisagov/pe-reports: A tool for creating Posture & Exposure reports.

Usage:
  pe-reports REPORT_DATE OUTPUT_DIRECTORY [--log-level=LEVEL]

Options:
  -h --help                         Show this message.
  REPORT_DATE                       Date of the report, format YYYY-MM-DD
  OUTPUT_DIRECTORY                  The directory where the final PDF
                                    reports should be saved.
  -l --log-level=LEVEL              If specified, then the log level will be set to
                                    the specified value.  Valid values are "debug", "info",
                                    "warning", "error", and "critical". [default: info]
"""

# Standard Python Libraries
import logging
import os
import sys
from typing import Any, Dict

# Third-Party Libraries
import docopt
import fitz
import pandas as pd
from schema import And, Schema, SchemaError, Use
from xhtml2pdf import pisa

from ._version import __version__
from .data.db_query import connect, get_orgs
from .pages import init


def embed(
    output_directory,
    org_code,
    datestring,
    file,
    cred_xlsx,
    da_xlsx,
    vuln_xlsx,
    mi_xlsx,
):
    """Embeds raw data into PDF and encrypts file."""
    doc = fitz.open(file)
    # Get the summary page of the PDF on page 4
    page = doc[3]
    output = (
        f"{output_directory}/{org_code}/Posture_and_Exposure_Report-{datestring}.pdf"
    )

    # Open CSV data as binary
    cc = open(cred_xlsx, "rb").read()
    da = open(da_xlsx, "rb").read()
    ma = open(vuln_xlsx, "rb").read()
    mi = open(mi_xlsx, "rb").read()

    # Insert link to CSV data in summary page of PDF.
    # Use coordinates to position them on the bottom.
    p1 = fitz.Point(110, 695)
    p2 = fitz.Point(240, 695)
    p3 = fitz.Point(375, 695)
    p5 = fitz.Point(500, 695)

    # Embed and add push-pin graphic
    page.add_file_annot(
        p1, cc, "compromised_credentials.xlsx", desc="Open up CSV", icon="PushPin"
    )
    page.add_file_annot(
        p2, da, "domain_alerts.xlsx", desc="Open up CSV", icon="PushPin"
    )
    page.add_file_annot(p3, ma, "vuln_alerts.xlsx", desc="Open up xlsx", icon="PushPin")
    page.add_file_annot(
        p5, mi, "mention_incidents.xlsx", desc="Open up CSV", icon="PushPin"
    )

    # Save doc and set garbage=4 to reduce PDF size using all 4 methods:
    # Remove unused objects, compact xref table, merge duplicate objects,
    # and check stream objects for duplication
    doc.save(
        output,
        garbage=4,
        deflate=True,
    )
    tooLarge = False
    # Throw error if file size is greater than 20MB
    filesize = os.path.getsize(output)
    if filesize >= 20000000:
        tooLarge = True

    return filesize, tooLarge


def convert_html_to_pdf(source_html, output_filename):
    """Convert HTML to PDF."""
    # Open output file for writing (truncated binary)
    result_file = open(output_filename, "w+b")

    # Convert HTML to PDF
    pisa_status = pisa.CreatePDF(
        source_html, dest=result_file  # the HTML to convert
    )  # file handle to receive result

    # Close output file
    result_file.close()  # close output file

    # Return False on success and True on errors
    return pisa_status.err


def generate_reports(datestring, output_directory):
    """Process steps for generating report data."""
    # Get PE orgs from PE db
    conn = connect()
    if conn:
        pe_orgs = get_orgs(conn)
    else:
        return 1
    generated_reports = 0

    # Iterate over organizations
    if pe_orgs:
        logging.info("PE orgs count: %d", len(pe_orgs))
        for org in pe_orgs:
            # Assign organization values
            org_uid = org[0]
            org_name = org[1]
            org_code = org[2]

            logging.info("Running on %s", org_code)

            # Create folders in output directory
            for dir_name in ("ppt", org_code):
                if not os.path.exists(f"{output_directory}/{dir_name}"):
                    os.mkdir(f"{output_directory}/{dir_name}")

            # Load source HTML
            try:
                basedir = os.path.abspath(os.path.dirname(__file__))
                template = os.path.join(basedir, "template.html")
                file = open(template)
                source_html = file.read().replace("\n", " ")
            except FileNotFoundError:
                logging.error(
                    "Template cannot be found. It must be named: '%s'", template
                )
                return 1

            # Insert Charts and Metrics into PDF
            (
                source_html,
                hibp_creds,
                cyber_creds,
                masq_df,
                insecure_df,
                vulns_df,
                assets_df,
                dark_web_mentions,
                alerts,
                top_cves,
            ) = init(
                source_html,
                datestring,
                org_name,
                org_uid,
            )

            # Close PDF
            file.close()

            # Convert to HTML to PDF
            output_filename = f"{output_directory}/{org_code}-Posture_and_Exposure_Report-{datestring}.pdf"
            convert_html_to_pdf(source_html, output_filename)

            # Create Credential Exposure Excel file
            cred_xlsx = f"{output_directory}/{org_code}/compromised_credentials.xlsx"
            credWriter = pd.ExcelWriter(cred_xlsx, engine="xlsxwriter")
            hibp_creds.to_excel(credWriter, sheet_name="HIBP_Credentials", index=False)
            cyber_creds.to_excel(
                credWriter, sheet_name="Cyber6_Credentials", index=False
            )
            credWriter.save()

            # Create Domain Masquerading Excel file
            da_xlsx = f"{output_directory}/{org_code}/domain_alerts.xlsx"
            domWriter = pd.ExcelWriter(da_xlsx, engine="xlsxwriter")
            masq_df.to_excel(domWriter, sheet_name="Suspected Domains", index=False)
            domWriter.save()

            # Create Suspected vulnerability Excel file
            vuln_xlsx = f"{output_directory}/{org_code}/vuln_alerts.xlsx"
            vulnWriter = pd.ExcelWriter(vuln_xlsx, engine="xlsxwriter")
            assets_df.to_excel(vulnWriter, sheet_name="Assets", index=False)
            insecure_df.to_excel(vulnWriter, sheet_name="Insecure", index=False)
            vulns_df.to_excel(vulnWriter, sheet_name="Verified Vulns", index=False)
            vulnWriter.save()

            # Create dark web Excel file
            mi_xlsx = f"{output_directory}/{org_code}/mention_incidents.xlsx"
            miWriter = pd.ExcelWriter(mi_xlsx, engine="xlsxwriter")
            dark_web_mentions.to_excel(
                miWriter, sheet_name="Dark Web Mentions", index=False
            )
            alerts.to_excel(miWriter, sheet_name="Dark Web Alerts", index=False)
            top_cves.to_excel(miWriter, sheet_name="Top CVEs", index=False)
            miWriter.save()

            # Grab the PDF
            pdf = f"{output_directory}/{org_code}-Posture_and_Exposure_Report-{datestring}.pdf"

            (filesize, tooLarge) = embed(
                output_directory,
                org_code,
                datestring,
                pdf,
                cred_xlsx,
                da_xlsx,
                vuln_xlsx,
                mi_xlsx,
            )
            # Log a message if the report is too large.  Our current mailer
            # cannot send files larger than 20MB.
            if tooLarge:
                logging.info(
                    "%s is too large. File size: %s Limit: 20MB", org_code, filesize
                )

            generated_reports += 1
    else:
        logging.error(
            "Connection to pe database failed and/or there are 0 organizations stored."
        )

    logging.info("%s reports generated", generated_reports)


def main():
    """Generate PDF reports."""
    args: Dict[str, str] = docopt.docopt(__doc__, version=__version__)

    # Validate and convert arguments as needed
    schema: Schema = Schema(
        {
            "--log-level": And(
                str,
                Use(str.lower),
                lambda n: n in ("debug", "info", "warning", "error", "critical"),
                error="Possible values for --log-level are "
                + "debug, info, warning, error, and critical.",
            ),
            str: object,  # Don't care about other keys, if any
        }
    )

    try:
        validated_args: Dict[str, Any] = schema.validate(args)
    except SchemaError as err:
        # Exit because one or more of the arguments were invalid
        print(err, file=sys.stderr)
        sys.exit(1)

    # Assign validated arguments to variables
    log_level: str = validated_args["--log-level"]

    # Set up logging
    logging.basicConfig(
        format="%(asctime)-15s %(levelname)s %(message)s", level=log_level.upper()
    )

    logging.info("Loading Posture & Exposure Report, Version : %s", __version__)

    # Create output directory
    if not os.path.exists(validated_args["OUTPUT_DIRECTORY"]):
        os.mkdir(validated_args["OUTPUT_DIRECTORY"])

    # Generate reports
    generate_reports(
        validated_args["REPORT_DATE"],
        validated_args["OUTPUT_DIRECTORY"],
    )

    # Stop logging and clean up
    logging.shutdown()

1	"""cisagov/pe-reports: A tool for creating Posture & Exposure reports.
2
3	Usage:
4	pe-reports REPORT_DATE OUTPUT_DIRECTORY [--log-level=LEVEL]
5
6	Options:
7	-h --help Show this message.
8	REPORT_DATE Date of the report, format YYYY-MM-DD
9	OUTPUT_DIRECTORY The directory where the final PDF
10	reports should be saved.
11	-l --log-level=LEVEL If specified, then the log level will be set to
12	the specified value. Valid values are "debug", "info",
13	"warning", "error", and "critical". [default: info]
14	"""
15
16	# Standard Python Libraries
17	import logging	5✔
18	import os	5✔
19	import sys	5✔
20	from typing import Any, Dict	5✔
21
22	# Third-Party Libraries
23	import docopt	5✔
24	import fitz	5✔
25	import pandas as pd	5✔
26	from schema import And, Schema, SchemaError, Use	5✔
27	from xhtml2pdf import pisa	5✔
28
29	from ._version import __version__	5✔
30	from .data.db_query import connect, get_orgs	5✔
31	from .pages import init	5✔
32
33
34	def embed(	5✔
35	output_directory,
36	org_code,
37	datestring,
38	file,
39	cred_xlsx,
40	da_xlsx,
41	vuln_xlsx,
42	mi_xlsx,
43	):
44	"""Embeds raw data into PDF and encrypts file."""
NEW 45	doc = fitz.open(file)	×
46	# Get the summary page of the PDF on page 4
NEW 47	page = doc[3]	×
NEW 48	output = (	×
49	f"{output_directory}/{org_code}/Posture_and_Exposure_Report-{datestring}.pdf"
50	)
51
52	# Open CSV data as binary
NEW 53	cc = open(cred_xlsx, "rb").read()	×
NEW 54	da = open(da_xlsx, "rb").read()	×
NEW 55	ma = open(vuln_xlsx, "rb").read()	×
NEW 56	mi = open(mi_xlsx, "rb").read()	×
57
58	# Insert link to CSV data in summary page of PDF.
59	# Use coordinates to position them on the bottom.
NEW 60	p1 = fitz.Point(110, 695)	×
NEW 61	p2 = fitz.Point(240, 695)	×
NEW 62	p3 = fitz.Point(375, 695)	×
NEW 63	p5 = fitz.Point(500, 695)	×
64
65	# Embed and add push-pin graphic
NEW 66	page.add_file_annot(	×
67	p1, cc, "compromised_credentials.xlsx", desc="Open up CSV", icon="PushPin"
68	)
NEW 69	page.add_file_annot(	×
70	p2, da, "domain_alerts.xlsx", desc="Open up CSV", icon="PushPin"
71	)
NEW 72	page.add_file_annot(p3, ma, "vuln_alerts.xlsx", desc="Open up xlsx", icon="PushPin")	×
NEW 73	page.add_file_annot(	×
74	p5, mi, "mention_incidents.xlsx", desc="Open up CSV", icon="PushPin"
75	)
76
77	# Save doc and set garbage=4 to reduce PDF size using all 4 methods:
78	# Remove unused objects, compact xref table, merge duplicate objects,
79	# and check stream objects for duplication
NEW 80	doc.save(	×
81	output,
82	garbage=4,
83	deflate=True,
84	)
NEW 85	tooLarge = False	×
86	# Throw error if file size is greater than 20MB
NEW 87	filesize = os.path.getsize(output)	×
NEW 88	if filesize >= 20000000:	×
NEW 89	tooLarge = True	×
90
NEW 91	return filesize, tooLarge	×
92
93
94	def convert_html_to_pdf(source_html, output_filename):	5✔
95	"""Convert HTML to PDF."""
96	# Open output file for writing (truncated binary)
NEW 97	result_file = open(output_filename, "w+b")	×
98
99	# Convert HTML to PDF
NEW 100	pisa_status = pisa.CreatePDF(	×
101	source_html, dest=result_file # the HTML to convert
102	) # file handle to receive result
103
104	# Close output file
NEW 105	result_file.close() # close output file	×
106
107	# Return False on success and True on errors
NEW 108	return pisa_status.err	×
109
110
111	def generate_reports(datestring, output_directory):	5✔
112	"""Process steps for generating report data."""
113	# Get PE orgs from PE db
114	conn = connect()	5✔
115	if conn:	5!
NEW 116	pe_orgs = get_orgs(conn)	×
117	else:
118	return 1	5✔
NEW 119	generated_reports = 0	×
120
121	# Iterate over organizations
NEW 122	if pe_orgs:	×
NEW 123	logging.info("PE orgs count: %d", len(pe_orgs))	×
NEW 124	for org in pe_orgs:	×
125	# Assign organization values
NEW 126	org_uid = org[0]	×
NEW 127	org_name = org[1]	×
NEW 128	org_code = org[2]	×
129
NEW 130	logging.info("Running on %s", org_code)	×
131
132	# Create folders in output directory
NEW 133	for dir_name in ("ppt", org_code):	×
NEW 134	if not os.path.exists(f"{output_directory}/{dir_name}"):	×
NEW 135	os.mkdir(f"{output_directory}/{dir_name}")	×
136
137	# Load source HTML
NEW 138	try:	×
NEW 139	basedir = os.path.abspath(os.path.dirname(__file__))	×
NEW 140	template = os.path.join(basedir, "template.html")	×
NEW 141	file = open(template)	×
NEW 142	source_html = file.read().replace("\n", " ")	×
NEW 143	except FileNotFoundError:	×
NEW 144	logging.error(	×
145	"Template cannot be found. It must be named: '%s'", template
146	)
NEW 147	return 1	×
148
149	# Insert Charts and Metrics into PDF
NEW 150	(	×
151	source_html,
152	hibp_creds,
153	cyber_creds,
154	masq_df,
155	insecure_df,
156	vulns_df,
157	assets_df,
158	dark_web_mentions,
159	alerts,
160	top_cves,
161	) = init(
162	source_html,
163	datestring,
164	org_name,
165	org_uid,
166	)
167
168	# Close PDF
NEW 169	file.close()	×
170
171	# Convert to HTML to PDF
NEW 172	output_filename = f"{output_directory}/{org_code}-Posture_and_Exposure_Report-{datestring}.pdf"	×
NEW 173	convert_html_to_pdf(source_html, output_filename)	×
174
175	# Create Credential Exposure Excel file
NEW 176	cred_xlsx = f"{output_directory}/{org_code}/compromised_credentials.xlsx"	×
NEW 177	credWriter = pd.ExcelWriter(cred_xlsx, engine="xlsxwriter")	×
NEW 178	hibp_creds.to_excel(credWriter, sheet_name="HIBP_Credentials", index=False)	×
NEW 179	cyber_creds.to_excel(	×
180	credWriter, sheet_name="Cyber6_Credentials", index=False
181	)
NEW 182	credWriter.save()	×
183
184	# Create Domain Masquerading Excel file
NEW 185	da_xlsx = f"{output_directory}/{org_code}/domain_alerts.xlsx"	×
NEW 186	domWriter = pd.ExcelWriter(da_xlsx, engine="xlsxwriter")	×
NEW 187	masq_df.to_excel(domWriter, sheet_name="Suspected Domains", index=False)	×
NEW 188	domWriter.save()	×
189
190	# Create Suspected vulnerability Excel file
NEW 191	vuln_xlsx = f"{output_directory}/{org_code}/vuln_alerts.xlsx"	×
NEW 192	vulnWriter = pd.ExcelWriter(vuln_xlsx, engine="xlsxwriter")	×
NEW 193	assets_df.to_excel(vulnWriter, sheet_name="Assets", index=False)	×
NEW 194	insecure_df.to_excel(vulnWriter, sheet_name="Insecure", index=False)	×
NEW 195	vulns_df.to_excel(vulnWriter, sheet_name="Verified Vulns", index=False)	×
NEW 196	vulnWriter.save()	×
197
198	# Create dark web Excel file
NEW 199	mi_xlsx = f"{output_directory}/{org_code}/mention_incidents.xlsx"	×
NEW 200	miWriter = pd.ExcelWriter(mi_xlsx, engine="xlsxwriter")	×
NEW 201	dark_web_mentions.to_excel(	×
202	miWriter, sheet_name="Dark Web Mentions", index=False
203	)
NEW 204	alerts.to_excel(miWriter, sheet_name="Dark Web Alerts", index=False)	×
NEW 205	top_cves.to_excel(miWriter, sheet_name="Top CVEs", index=False)	×
NEW 206	miWriter.save()	×
207
208	# Grab the PDF
NEW 209	pdf = f"{output_directory}/{org_code}-Posture_and_Exposure_Report-{datestring}.pdf"	×
210
NEW 211	(filesize, tooLarge) = embed(	×
212	output_directory,
213	org_code,
214	datestring,
215	pdf,
216	cred_xlsx,
217	da_xlsx,
218	vuln_xlsx,
219	mi_xlsx,
220	)
221	# Log a message if the report is too large. Our current mailer
222	# cannot send files larger than 20MB.
NEW 223	if tooLarge:	×
NEW 224	logging.info(	×
225	"%s is too large. File size: %s Limit: 20MB", org_code, filesize
226	)
227
NEW 228	generated_reports += 1	×
229	else:
NEW 230	logging.error(	×
231	"Connection to pe database failed and/or there are 0 organizations stored."
232	)
233
NEW 234	logging.info("%s reports generated", generated_reports)	×
235
236
237	def main():	5✔
238	"""Generate PDF reports."""
239	args: Dict[str, str] = docopt.docopt(__doc__, version=__version__)	5✔
240
241	# Validate and convert arguments as needed
242	schema: Schema = Schema(	5✔
243	{
244	"--log-level": And(
245	str,
246	Use(str.lower),
247	lambda n: n in ("debug", "info", "warning", "error", "critical"),
248	error="Possible values for --log-level are "
249	+ "debug, info, warning, error, and critical.",
250	),
251	str: object, # Don't care about other keys, if any
252	}
253	)
254
255	try:	5✔
256	validated_args: Dict[str, Any] = schema.validate(args)	5✔
257	except SchemaError as err:	5✔
258	# Exit because one or more of the arguments were invalid
259	print(err, file=sys.stderr)	5✔
260	sys.exit(1)	5✔
261
262	# Assign validated arguments to variables
263	log_level: str = validated_args["--log-level"]	5✔
264
265	# Set up logging
266	logging.basicConfig(	5✔
267	format="%(asctime)-15s %(levelname)s %(message)s", level=log_level.upper()
268	)
269
270	logging.info("Loading Posture & Exposure Report, Version : %s", __version__)	5✔
271
272	# Create output directory
273	if not os.path.exists(validated_args["OUTPUT_DIRECTORY"]):	5✔
274	os.mkdir(validated_args["OUTPUT_DIRECTORY"])	5✔
275
276	# Generate reports
277	generate_reports(	5✔
278	validated_args["REPORT_DATE"],
279	validated_args["OUTPUT_DIRECTORY"],
280	)
281
282	# Stop logging and clean up
283	logging.shutdown()	5✔

cisagov / pe-reports / 1773811980

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous