add export discovery assets to CSV format

Author: suskind

discovery_assets_to_csv.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python
+"""
+Export all discovery assets with a specific optional score to CSV format
+
+Run: 
+
+$ python3 discovery_assets_to_csv.py -s <OPTIONAL_SCORE> -o <OUTPUT_FILE_PATH>
+
+"""
+import argparse
+import requests
+import csv
+from urllib.parse import urljoin, quote
+from datetime import datetime
+
+# Define the JWT or it will be asked when you run the script
+jwt_token = None
+
+api_base_url = 'https://api.probely.com'
+discovery_assets_endpoint = urljoin(api_base_url, "discovery/assets/?length=10000&page=1&ordering=-last_seen&{score_str}")
+
+def map_risk(probely_risk):
+    if probely_risk == 10:
+        return 'Low'
+    elif probely_risk == 20:
+        return 'Normal'
+    elif probely_risk == 30:
+        return 'High'
+    else:
+        return 'NA'
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-s', '--score', help='Score', required=False, choices=[
+        'A+', 'A', 'B', 'C', 'D', 'E', 'F', 'R', 'NA'])
+    parser.add_argument('-o', '--output', help='Output CSV file', type=argparse.FileType('w'), required=True)
+    args = parser.parse_args()
+
+    if jwt_token is None:
+        token = input("API Token:")
+    else:
+        token = jwt_token
+
+    if token is None or token == '':
+        print('Error: JWT is required')
+        return
+    headers = {'Authorization': "JWT {}".format(token)}
+
+    score_str = f'score={quote(args.score)}'
+    if args.score is None:
+        score_str = ''
+
+    response_assets = requests.get(
+        discovery_assets_endpoint.format(score_str=score_str),
+        headers=headers
+    )
+    response_assets.raise_for_status()
+    assets_res = response_assets.json()['results']
+
+    csv_writer = csv.writer(
+            args.output, delimiter=",", quotechar='"', quoting=csv.QUOTE_ALL
+        )
+    row = ['ID', 'Name', 'URL', 'Type', 'Last Seen', 'Risk', 'Score', 'State']
+    csv_writer.writerow(row)
+    for asset in assets_res:
+        row = [
+            asset['id'],
+            asset['name'],
+            asset['url'],
+            asset['type'],
+            datetime.strptime(asset['last_seen'], "%Y-%m-%dT%H:%M:%SZ").strftime("%Y-%m-%d %H:%M:%S"),
+            map_risk(asset['risk']),
+            asset['score'],
+            asset['state'],
+        ]
+        csv_writer.writerow(row)
+
+    print('Done')
+
+if __name__ == '__main__':
+    main()