Generalize protocol v1 proxy client actions

ocf/loxthelion/managed-container.d/proxy_client.py

@@ -9,18 +9,18 @@ PROTOCOL_VERSION = 1
 VALID_OCF_RCS = {0, 1, 7, 8}
 
 
-def main():
-    if len(sys.argv) != 5:
-        raise SystemExit(11)
+class TransportFailure(Exception):
+    pass
 
-    socket_path = sys.argv[1]
-    resource = sys.argv[2]
-    transport_failure = int(sys.argv[3])
-    protocol_failure = int(sys.argv[4])
 
+class ProtocolFailure(Exception):
+    pass
+
+
+def send_request(socket_path, resource, action):
     request = {
         "protocol_version": PROTOCOL_VERSION,
-        "action": "monitor",
+        "action": action,
         "resource": resource,
     }
 
@@ -35,7 +35,10 @@ def main():
             sock.connect(socket_path)
 
             sock.sendall(
-                json.dumps(request).encode("utf-8")
+                json.dumps(
+                    request,
+                    separators=(",", ":"),
+                ).encode("utf-8")
                 + b"\n"
             )
 
@@ -56,49 +59,60 @@ def main():
         socket.timeout,
         ConnectionError,
         OSError,
-    ):
-        raise SystemExit(transport_failure)
+    ) as exc:
+        raise TransportFailure from exc
 
-    # A connection was established. From this point onward, absence of a
-    # valid response is a protocol/result failure rather than transport loss.
+    return raw
+
+
+def decode_response(raw):
     if not raw:
-        raise SystemExit(protocol_failure)
+        raise ProtocolFailure
 
     try:
-        obj = json.loads(
+        response = json.loads(
             raw.decode("utf-8").strip()
         )
 
     except (
         UnicodeDecodeError,
         json.JSONDecodeError,
-    ):
-        raise SystemExit(protocol_failure)
+    ) as exc:
+        raise ProtocolFailure from exc
 
-    if not isinstance(obj, dict):
-        raise SystemExit(protocol_failure)
+    if not isinstance(response, dict):
+        raise ProtocolFailure
 
-    if obj.get("protocol_version") != PROTOCOL_VERSION:
-        raise SystemExit(protocol_failure)
+    return response
 
-    if obj.get("ok") is not True:
-        raise SystemExit(protocol_failure)
 
-    if obj.get("resource") != resource:
-        raise SystemExit(protocol_failure)
+def validate_response(response, resource, action):
+    if response.get("protocol_version") != PROTOCOL_VERSION:
+        raise ProtocolFailure
 
-    if obj.get("action") != "monitor":
-        raise SystemExit(protocol_failure)
+    if response.get("ok") is not True:
+        raise ProtocolFailure
 
-    result = obj.get("result")
+    if response.get("resource") != resource:
+        raise ProtocolFailure
 
+    if response.get("action") != action:
+        raise ProtocolFailure
+
+    if "result" not in response:
+        raise ProtocolFailure
+
+    return response["result"]
+
+
+def format_monitor_result(result):
     if not isinstance(result, dict):
-        raise SystemExit(protocol_failure)
+        raise ProtocolFailure
 
     ha = result.get("ha")
 
     if not isinstance(ha, dict):
-        raise SystemExit(protocol_failure)
+        raise ProtocolFailure
 
     try:
         ocf_rc = int(ha.get("ocf_rc"))
@@ -106,18 +120,18 @@ def main():
     except (
         TypeError,
         ValueError,
-    ):
-        raise SystemExit(protocol_failure)
+    ) as exc:
+        raise ProtocolFailure from exc
 
     if ocf_rc not in VALID_OCF_RCS:
-        raise SystemExit(protocol_failure)
+        raise ProtocolFailure
 
     receiver = result.get("wal_receiver") or {}
 
     if not isinstance(receiver, dict):
-        raise SystemExit(protocol_failure)
+        raise ProtocolFailure
 
-    payload = "|".join([
+    return "|".join([
         str(ocf_rc),
         str(ha.get("health", "unknown")),
         str(ha.get("ha_role", "unknown")),
@@ -139,7 +153,62 @@ def main():
         str(receiver.get("sender_host") or ""),
     ])
 
-    print(payload)
+
+def format_result(action, result):
+    if action == "monitor":
+        return format_monitor_result(result)
+
+    return json.dumps(
+        result,
+        separators=(",", ":"),
+        sort_keys=True,
+    )
+
+
+def main():
+    if len(sys.argv) != 6:
+        raise SystemExit(11)
+
+    socket_path = sys.argv[1]
+    resource = sys.argv[2]
+    action = sys.argv[3]
+
+    try:
+        transport_failure = int(sys.argv[4])
+        protocol_failure = int(sys.argv[5])
+
+    except ValueError:
+        raise SystemExit(11)
+
+    if not resource or not action:
+        raise SystemExit(protocol_failure)
+
+    try:
+        raw = send_request(
+            socket_path,
+            resource,
+            action,
+        )
+
+    except TransportFailure:
+        raise SystemExit(transport_failure)
+
+    try:
+        response = decode_response(raw)
+        result = validate_response(
+            response,
+            resource,
+            action,
+        )
+        output = format_result(
+            action,
+            result,
+        )
+
+    except ProtocolFailure:
+        raise SystemExit(protocol_failure)
+
+    print(output)
 
 
 if __name__ == "__main__":