From fe141170d2c13cdbdf369cd7fbb0456891acb3a3 Mon Sep 17 00:00:00 2001
From: Bridget <bridget.littlebird@pm.me>
Date: Sun, 12 Apr 2020 18:27:18 +0200
Subject: [PATCH] Refactor

---
 lslsh.py | 61 +++++++++++++++++++++++++++++++-------------------------
 1 file changed, 34 insertions(+), 27 deletions(-)

diff --git a/lslsh.py b/lslsh.py
index c3d242b..432802f 100644
--- a/lslsh.py
+++ b/lslsh.py
@@ -3,35 +3,46 @@ import warnings
 from urllib3.connectionpool import InsecureRequestWarning
 from json.decoder import JSONDecodeError
 import sys
+from typing import Dict
 
 SECRET_KEY: str = "29731e5170353a8b235098c43cd2099a4e805c55fb4395890e81f437c17334a9"
 
 warnings.filterwarnings("ignore", category=InsecureRequestWarning)
 
+def send_cmd(url: str, cmd: str) -> Dict:
+    data = {"secret_key": SECRET_KEY,
+            "command": cmd}
+
+    try:
+        response = requests.post(url, json=data, verify=False)
+        response_data = response.json()
+    except (requests.ConnectionError, requests.exceptions.MissingSchema,
+            requests.exceptions.InvalidURL):
+        raise requests.exceptions.InvalidURL
+    except JSONDecodeError:
+        raise Exception("Error: Response has malformed json")
+
+    error = response_data.get("error", None)
+    if error:
+        raise Exception(f"Error: {error}")
+
+    return response_data
+
 def init() -> str:
     initialized: bool = False
     while not initialized:
+        url = input("url > ")
         try:
-            data = {"secret_key": SECRET_KEY,
-                    "command": "init"}
-            url = input("url > ")
-            response = requests.post(url, json=data, verify=False)
-            error = response.json().get("error")
-            if error:
-                print(f"Error: {error}")
-                sys.exit(1)
-        except (requests.ConnectionError, requests.exceptions.MissingSchema,
-                requests.exceptions.InvalidURL):
+            result = send_cmd(url, "init")
+        except requests.exceptions.InvalidURL:
             print("Error: Invalid URL")
-        except (KeyError, JSONDecodeError):
-            print(f"Error: Invalid response")
-        else:
-            try:
-                uuid = response.json().get("uuid")
-            except KeyError:
-                print(f"Error: Invalid response")
-            else:
-                initialized = True
+            continue
+
+        uuid = result.get("uuid", None)
+        if not uuid:
+            raise Exception("Error: Invalid response")
+
+        initialized = True
 
     print(f"Connected to {uuid}")
     print("_______________________________________________________________________________")
@@ -39,21 +50,17 @@ def init() -> str:
     return url
 
 def disconnect(url: str):
-    data = {"secret_key": SECRET_KEY,
-            "command": "disconnect"}
-    response = requests.post(url, json=data, verify=False)
+    result = send_cmd(url, "disconnect")
     print("")
-    if response.content.decode("utf-8") == "disconnected":
+    if result.get("result") == "disconnected":
         print("Disconnected from remote.")
     else:
         print("Disconnected from remote (without acknowledgement).")
 
 def run(url: str):
     while True:
-        data = {"secret_key": SECRET_KEY,
-                "command": input("sl > ")}
-        response = requests.post(url, json=data, verify=False)
-        print(response.content.decode("utf-8"))
+        result = send_cmd(url, input("sl > "))
+        print(result.get("result"))
 
 try:
     url = init()