From 0975315801fec1c2ceb42dae0d040aab7ed9bb92 Mon Sep 17 00:00:00 2001
From: Guillaume <65369351+guillaumeVilar@users.noreply.github.com>
Date: Fri, 8 Nov 2024 10:20:14 +0000
Subject: [PATCH] examples: improve get_token to be able to pull only the ssl
 cert (#29)

---
 examples/get_token.py | 41 ++++++++++++++++++++++++++++++-----------
 1 file changed, 30 insertions(+), 11 deletions(-)

diff --git a/examples/get_token.py b/examples/get_token.py
index c02f9458..a246ca97 100755
--- a/examples/get_token.py
+++ b/examples/get_token.py
@@ -4,31 +4,50 @@
 # that can be found in the COPYING file.
 
 
-# Fetches a session token and optional SSL certificate from CVP.
-# Useful when authenticating in the included examples.
+# Fetches a session token and/or the SSL certificate from CVP.
+# Helpful for authentication in the provided examples.
 
 import argparse
 import requests
 import ssl
 import json
 
+TOKEN_FILE = 'token.txt'
+SSL_CERT_FILE = 'cvp.crt'
+
 
 def main(args):
-    r = requests.post('https://' + args.server + '/cvpservice/login/authenticate.do',
-                      auth=(args.username, args.password), verify=args.ssl is False)
+    if (not args.username and not args.password) and not args.ssl:
+        exit(
+            "Error: Arguments --username and --password are required (for token generation), "
+            "and/or --ssl argument is required (for SSL cert generation)."
+        )
+        return
+
+    if args.username and args.password:
+        r = requests.post(
+            'https://' + args.server + '/cvpservice/login/authenticate.do',
+            auth=(args.username, args.password), verify=args.ssl is False)
 
-    r.json()['sessionId']
+        r.json()['sessionId']
 
-    with open("token.txt", "w") as f:
-        f.write(r.json()['sessionId'])
+        with open(TOKEN_FILE, "w") as f:
+            f.write(r.json()['sessionId'])
 
     if args.ssl:
-        with open("cvp.crt", "w") as f:
+        with open(SSL_CERT_FILE, "w") as f:
             f.write(ssl.get_server_certificate((args.server, 443)))
 
 
 if __name__ == '__main__':
-    ds = ("Get a session token (and optional SSL cert) from CVP and store to token.txt")
+    ds = (
+        "Get a session token and/or the SSL cert from CVP, and store them respectively "
+        "in token.txt and cvp.crt.\n"
+        "Example 1 (token.txt and cvp.crt generation): \n  "
+        "python3 get_token.py --server 10.0.0.1 --username <username> --password <password> --ssl\n"
+        "Example 2 (cvp.crt generation only): \n  "
+        "python3 get_token.py --server 10.0.0.1 --ssl"
+    )
     parser = argparse.ArgumentParser(
         description=ds,
         formatter_class=argparse.RawDescriptionHelpFormatter)
@@ -36,9 +55,9 @@ def main(args):
         '--server',
         required=True,
         help="CloudVision server to connect to in <host> format.")
-    parser.add_argument("--username", required=True, type=str,
+    parser.add_argument("--username", type=str,
                         help="Username to authorize with")
-    parser.add_argument("--password", required=True, type=str,
+    parser.add_argument("--password", type=str,
                         help="Password to authorize with")
     parser.add_argument("--ssl", action="store_true",
                         help="Save the self-signed certficate to cvp.crt")