From 3a2c0ec3924541f647c8fec6c36d25f8c3ec3221 Mon Sep 17 00:00:00 2001 From: 0x4A616B65 Date: Fri, 9 Aug 2024 09:29:28 -0400 Subject: [PATCH 1/3] Added override_import_dt to sql --- docs/sql.md | 16 +++++++++++++++- qualyspy/sql/base.py | 5 +++-- qualyspy/sql/supported_uploads.py | 26 ++++++++++++++++---------- 3 files changed, 34 insertions(+), 13 deletions(-) diff --git a/docs/sql.md b/docs/sql.md index 2c44a04..1a4bf1c 100644 --- a/docs/sql.md +++ b/docs/sql.md @@ -44,7 +44,10 @@ And finally, you can use the following supported functions: >**Head's Up:** More upload functions are coming soon! -Each upload function takes 2 parameters. The first is the ```BaseList``` of data, and the second is the ```sqlalchemy.Connection``` object you built above. +Each upload function takes 2 positional parameters. The first is the ```BaseList``` of data, and the second is the ```sqlalchemy.Connection``` object you built above. + +Functions also take an optional ```override_import_dt``` parameter that will set the resulting SQL table's ```import_datetime``` field to the value you specify. ```override_import_dt``` is a ```datetime.datetime``` object. + | Function Name | Module | ```qualyspy``` Function Data Source | Resulting Table Name | | -- | -- | -- | -- | @@ -54,6 +57,17 @@ Each upload function takes 2 parameters. The first is the ```BaseList``` of data | ```upload_vmdr_hld``` | VMDR | ```vmdr.get_hld()```| ```vmdr_hld_hosts_list``` for hosts and ```vmdr_hld_detections``` for detections | | ```upload_vmdr_ips``` | VMDR | ```vmdr.get_ip_list()```| ```vmdr_ips``` | +```py +from qualyspy.sql import * + +# Get a connection to the DB +cnxn = db_connect(host='10.0.0.1', db='qualysdata', trusted_cnxn=True) + +# Upload a previous vmdr.get_host_list() call to the DB, with override_import_dt set +# to 10-25-2023 12:00:00 +uploaded = upload_vmdr_hosts(vmdr_hosts, cnxn, override_import_dt=datetime.datetime(2023, 10, 25, 12, 0, 0)) +>>>Uploaded 12345 records to vmdr_hosts_list + ## A Friendly Recommendation For Getting Data When calling any of the data source functions to get the data to upload, it is recommended to make the call as verbose as possible via kwargs. diff --git a/qualyspy/sql/base.py b/qualyspy/sql/base.py index a88c500..bf3b08f 100644 --- a/qualyspy/sql/base.py +++ b/qualyspy/sql/base.py @@ -75,7 +75,7 @@ def db_connect( return engine.connect() -def upload_data(df: DataFrame, table: str, cnxn: Connection, dtype: dict) -> int: +def upload_data(df: DataFrame, table: str, cnxn: Connection, dtype: dict, override_import_dt: datetime = None) -> int: """ Upload a DataFrame to a SQL table. Appends 'import_datetime' column to the DataFrame. @@ -84,13 +84,14 @@ def upload_data(df: DataFrame, table: str, cnxn: Connection, dtype: dict) -> int table (str): The name of the table to upload to. cnxn (Connection): The Connection object to the SQL database. dtype (dict): The data types of the columns in the table. Key is the column name, value is the data type as sqlalchemy.types.Something() + override_import_dt (datetime): If provided, will override the import_datetime column with this value. Returns: int: The number of rows uploaded. """ # Add an import_datetime column: - df["import_datetime"] = datetime.now() + df["import_datetime"] = datetime.now() if not override_import_dt else override_import_dt dtype["import_datetime"] = types.DateTime() # For any string values in the DataFrame, make sure it doesn't diff --git a/qualyspy/sql/supported_uploads.py b/qualyspy/sql/supported_uploads.py index 232bfd6..e95dee7 100644 --- a/qualyspy/sql/supported_uploads.py +++ b/qualyspy/sql/supported_uploads.py @@ -2,6 +2,8 @@ supported_uploads.py - Contains the functions to upload supported API pulls to SQL Server. """ +from datetime import datetime + from pandas import DataFrame from sqlalchemy import Connection, types @@ -9,13 +11,14 @@ from ..vmdr.data_classes.lists import BaseList -def upload_vmdr_ags(ags: BaseList, cnxn: Connection) -> int: +def upload_vmdr_ags(ags: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int: """ Upload data from vmdr.get_ag_list() to SQL. Parameters: ags (BaseList): The Asset Group List to upload. cnxn (Connection): The Connection object to the SQL database. + override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database. Returns: int: The number of rows uploaded. @@ -63,16 +66,17 @@ def upload_vmdr_ags(ags: BaseList, cnxn: Connection) -> int: df = DataFrame([prepare_dataclass(ag) for ag in ags]) # Upload the data: - return upload_data(df, "vmdr_assetgroups", cnxn, dtype=COLS) + return upload_data(df, "vmdr_assetgroups", cnxn, dtype=COLS, override_import_dt=override_import_dt) -def upload_vmdr_kb(kbs: BaseList, cnxn: Connection) -> int: +def upload_vmdr_kb(kbs: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int: """ Upload data from vmdr.query_kb() to SQL. Parameters: kbs (BaseList): The KB List to upload. cnxn (Connection): The Connection object to the SQL database. + override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database. Returns: int: The number of rows uploaded. @@ -115,10 +119,10 @@ def upload_vmdr_kb(kbs: BaseList, cnxn: Connection) -> int: df = DataFrame([prepare_dataclass(kb) for kb in kbs]) # Upload the data: - return upload_data(df, "vmdr_knowledgebase", cnxn, dtype=COLS) + return upload_data(df, "vmdr_knowledgebase", cnxn, dtype=COLS, override_import_dt=override_import_dt) -def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False) -> int: +def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False, override_import_dt: datetime = None) -> int: """ Upload data from vmdr.get_host_list() to SQL. @@ -126,6 +130,7 @@ def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False) - hosts (BaseList): The Host List to upload. cnxn (Connection): The Connection object to the SQL database. is_hld (bool): If the data is from a Host List Detail pull. You can ignore this. + override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database. Returns: int: The number of rows uploaded. @@ -192,17 +197,18 @@ def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False) - # Upload the data, with table depdening on if it is a Host List Detail pull or not: return upload_data( - df, "vmdr_hosts_list" if not is_hld else "vmdr_hld_hosts_list", cnxn, dtype=COLS + df, "vmdr_hosts_list" if not is_hld else "vmdr_hld_hosts_list", cnxn, dtype=COLS, override_import_dt=override_import_dt ) -def upload_vmdr_ips(ips: BaseList, cnxn: Connection) -> int: +def upload_vmdr_ips(ips: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int: """ Upload data from vmdr.get_ip_list() to SQL. Parameters: ips (BaseList): The IP List to upload from vmdr.get_ip_list(). cnxn (Connection): The Connection object to the SQL database. + override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database. Returns: int: The number of rows uploaded. @@ -220,10 +226,10 @@ def upload_vmdr_ips(ips: BaseList, cnxn: Connection) -> int: df["TYPE"] = df["IP"].apply(lambda x: "Single IP" if "/" not in x else "IP Range") # Upload the data: - return upload_data(df, "vmdr_ips", cnxn, dtype=COLS) + return upload_data(df, "vmdr_ips", cnxn, dtype=COLS, override_import_dt=override_import_dt) -def upload_vmdr_hld(hld: BaseList, cnxn: Connection) -> int: +def upload_vmdr_hld(hld: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int: """ Upload data from vmdr.get_hld() to SQL. @@ -290,4 +296,4 @@ def upload_vmdr_hld(hld: BaseList, cnxn: Connection) -> int: df["QDS"] = df["QDS"].apply(lambda x: int(x) if x else None) # Upload the data: - return upload_data(df, "vmdr_hld_detections", cnxn, dtype=COLS) + return upload_data(df, "vmdr_hld_detections", cnxn, dtype=COLS, override_import_dt=override_import_dt) From 14ed34fddbc88df49f70f8beed42b0690108c06a Mon Sep 17 00:00:00 2001 From: 0x4A616B65 Date: Fri, 9 Aug 2024 09:30:16 -0400 Subject: [PATCH 2/3] Added override_import_dt to sql --- docs/sql.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/sql.md b/docs/sql.md index 1a4bf1c..fa1b5aa 100644 --- a/docs/sql.md +++ b/docs/sql.md @@ -65,7 +65,8 @@ cnxn = db_connect(host='10.0.0.1', db='qualysdata', trusted_cnxn=True) # Upload a previous vmdr.get_host_list() call to the DB, with override_import_dt set # to 10-25-2023 12:00:00 -uploaded = upload_vmdr_hosts(vmdr_hosts, cnxn, override_import_dt=datetime.datetime(2023, 10, 25, 12, 0, 0)) +dt = datetime.datetime(2023, 10, 25, 12, 0, 0) +uploaded = upload_vmdr_hosts(vmdr_hosts, cnxn, override_import_dt=dt) >>>Uploaded 12345 records to vmdr_hosts_list ## A Friendly Recommendation For Getting Data From 2d3cfa9ceb7df2dac42477d9505e8baf6cbc170d Mon Sep 17 00:00:00 2001 From: jake-lindsay-tfs <121048664+jake-lindsay-tfs@users.noreply.github.com> Date: Fri, 9 Aug 2024 13:30:47 +0000 Subject: [PATCH 3/3] :art: Format Python code with psf/black --- qualyspy/sql/base.py | 12 +++++-- qualyspy/sql/supported_uploads.py | 55 ++++++++++++++++++++++++------- 2 files changed, 54 insertions(+), 13 deletions(-) diff --git a/qualyspy/sql/base.py b/qualyspy/sql/base.py index bf3b08f..3d2f1db 100644 --- a/qualyspy/sql/base.py +++ b/qualyspy/sql/base.py @@ -75,7 +75,13 @@ def db_connect( return engine.connect() -def upload_data(df: DataFrame, table: str, cnxn: Connection, dtype: dict, override_import_dt: datetime = None) -> int: +def upload_data( + df: DataFrame, + table: str, + cnxn: Connection, + dtype: dict, + override_import_dt: datetime = None, +) -> int: """ Upload a DataFrame to a SQL table. Appends 'import_datetime' column to the DataFrame. @@ -91,7 +97,9 @@ def upload_data(df: DataFrame, table: str, cnxn: Connection, dtype: dict, overri """ # Add an import_datetime column: - df["import_datetime"] = datetime.now() if not override_import_dt else override_import_dt + df["import_datetime"] = ( + datetime.now() if not override_import_dt else override_import_dt + ) dtype["import_datetime"] = types.DateTime() # For any string values in the DataFrame, make sure it doesn't diff --git a/qualyspy/sql/supported_uploads.py b/qualyspy/sql/supported_uploads.py index e95dee7..b1f1061 100644 --- a/qualyspy/sql/supported_uploads.py +++ b/qualyspy/sql/supported_uploads.py @@ -11,14 +11,16 @@ from ..vmdr.data_classes.lists import BaseList -def upload_vmdr_ags(ags: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int: +def upload_vmdr_ags( + ags: BaseList, cnxn: Connection, override_import_dt: datetime = None +) -> int: """ Upload data from vmdr.get_ag_list() to SQL. Parameters: ags (BaseList): The Asset Group List to upload. cnxn (Connection): The Connection object to the SQL database. - override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database. + override_import_dt (datetime): Use the passed datetime instead of generating one to upload to the database. Returns: int: The number of rows uploaded. @@ -66,10 +68,14 @@ def upload_vmdr_ags(ags: BaseList, cnxn: Connection, override_import_dt: datetim df = DataFrame([prepare_dataclass(ag) for ag in ags]) # Upload the data: - return upload_data(df, "vmdr_assetgroups", cnxn, dtype=COLS, override_import_dt=override_import_dt) + return upload_data( + df, "vmdr_assetgroups", cnxn, dtype=COLS, override_import_dt=override_import_dt + ) -def upload_vmdr_kb(kbs: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int: +def upload_vmdr_kb( + kbs: BaseList, cnxn: Connection, override_import_dt: datetime = None +) -> int: """ Upload data from vmdr.query_kb() to SQL. @@ -119,10 +125,21 @@ def upload_vmdr_kb(kbs: BaseList, cnxn: Connection, override_import_dt: datetime df = DataFrame([prepare_dataclass(kb) for kb in kbs]) # Upload the data: - return upload_data(df, "vmdr_knowledgebase", cnxn, dtype=COLS, override_import_dt=override_import_dt) + return upload_data( + df, + "vmdr_knowledgebase", + cnxn, + dtype=COLS, + override_import_dt=override_import_dt, + ) -def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False, override_import_dt: datetime = None) -> int: +def upload_vmdr_hosts( + hosts: BaseList, + cnxn: Connection, + is_hld: bool = False, + override_import_dt: datetime = None, +) -> int: """ Upload data from vmdr.get_host_list() to SQL. @@ -197,11 +214,17 @@ def upload_vmdr_hosts(hosts: BaseList, cnxn: Connection, is_hld: bool = False, o # Upload the data, with table depdening on if it is a Host List Detail pull or not: return upload_data( - df, "vmdr_hosts_list" if not is_hld else "vmdr_hld_hosts_list", cnxn, dtype=COLS, override_import_dt=override_import_dt + df, + "vmdr_hosts_list" if not is_hld else "vmdr_hld_hosts_list", + cnxn, + dtype=COLS, + override_import_dt=override_import_dt, ) -def upload_vmdr_ips(ips: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int: +def upload_vmdr_ips( + ips: BaseList, cnxn: Connection, override_import_dt: datetime = None +) -> int: """ Upload data from vmdr.get_ip_list() to SQL. @@ -226,10 +249,14 @@ def upload_vmdr_ips(ips: BaseList, cnxn: Connection, override_import_dt: datetim df["TYPE"] = df["IP"].apply(lambda x: "Single IP" if "/" not in x else "IP Range") # Upload the data: - return upload_data(df, "vmdr_ips", cnxn, dtype=COLS, override_import_dt=override_import_dt) + return upload_data( + df, "vmdr_ips", cnxn, dtype=COLS, override_import_dt=override_import_dt + ) -def upload_vmdr_hld(hld: BaseList, cnxn: Connection, override_import_dt: datetime = None) -> int: +def upload_vmdr_hld( + hld: BaseList, cnxn: Connection, override_import_dt: datetime = None +) -> int: """ Upload data from vmdr.get_hld() to SQL. @@ -296,4 +323,10 @@ def upload_vmdr_hld(hld: BaseList, cnxn: Connection, override_import_dt: datetim df["QDS"] = df["QDS"].apply(lambda x: int(x) if x else None) # Upload the data: - return upload_data(df, "vmdr_hld_detections", cnxn, dtype=COLS, override_import_dt=override_import_dt) + return upload_data( + df, + "vmdr_hld_detections", + cnxn, + dtype=COLS, + override_import_dt=override_import_dt, + )