From e3c2c1565d0eb6e71e3cfd4adfeb017ea35c81cd Mon Sep 17 00:00:00 2001
From: Daniel Micay <daniel.micay@grapheneos.org>
Date: Fri, 5 Jul 2024 00:35:51 -0400
Subject: [PATCH] ovh-mitigation: add checking/toggling firewall

---
 ovh-mitigation.py | 67 +++++++++++++++++++++++++++++------------------
 1 file changed, 41 insertions(+), 26 deletions(-)

diff --git a/ovh-mitigation.py b/ovh-mitigation.py
index 81d7764..f8c7fe1 100644
--- a/ovh-mitigation.py
+++ b/ovh-mitigation.py
@@ -8,14 +8,15 @@ def err(*args, **kwargs):
     print(*args, file=sys.stderr, **kwargs)
     sys.exit(1)
 
-if len(sys.argv) != 3:
-    err("expected 2 arguments")
+if len(sys.argv) != 4:
+    err("expected 3 arguments")
 
 with open("ovh-mitigation.json") as config:
     config = json.load(config)
 
-action = sys.argv[1]
-ip = sys.argv[2]
+system = sys.argv[1]
+action = sys.argv[2]
+ip = sys.argv[3]
 
 client = ovh.Client(
     endpoint=config["endpoint"],
@@ -24,28 +25,42 @@ client = ovh.Client(
     consumer_key=config["consumer_key"],
 )
 
-if action == "status":
-    result = client.get(f"/ip/{ip}/mitigation")
-    if len(result) == 0:
-        print("mitigation disabled")
-    elif len(result) == 1:
-        result = client.get(f"/ip/{ip}/mitigation/{result[0]}")
-        print("automatic: " + str(result["auto"]))
-        print("permanent: " + str(result["permanent"]))
+if system == "mitigation":
+    if action == "status":
+        result = client.get(f"/ip/{ip}/mitigation")
+        if len(result) == 0:
+            print("mitigation disabled")
+        elif len(result) == 1:
+            result = client.get(f"/ip/{ip}/mitigation/{result[0]}")
+            print("automatic: " + str(result["auto"]))
+            print("permanent: " + str(result["permanent"]))
+            print("state: " + result["state"])
+        else:
+            err("expected single result")
+    elif action == "enable":
+        result = client.post(f'/ip/{ip}/mitigation', ipOnMitigation=ip)
+        if result["permanent"] is True and result["state"] == "creationPending":
+            print("enabling permanent mitigation")
+        else:
+            err(json.dumps(result, indent=4) + "\nunexpected result")
+    elif action == "disable":
+        result = client.delete(f"/ip/{ip}/mitigation/{ip}")
+        if result["permanent"] is True and result["state"] == "removalPending":
+            print("disabling permanent mitigation")
+        else:
+            err(json.dumps(result, indent=4) + "\nunexpected result")
+    else:
+        err("unknown action: " + action)
+elif system == "firewall":
+    if action == "status":
+        result = client.get(f"/ip/{ip}/firewall/{ip}")
+        print("enabled: " + str(result["enabled"]))
         print("state: " + result["state"])
+    elif action == "enable":
+        client.put(f'/ip/{ip}/firewall/{ip}', enabled=True)
+    elif action == "disable":
+        client.put(f"/ip/{ip}/firewall/{ip}", enabled=False)
     else:
-        err("expected single result")
-elif action == "enable":
-    result = client.post(f'/ip/{ip}/mitigation', ipOnMitigation=ip)
-    if result["permanent"] is True and result["state"] == "creationPending":
-        print("enabling permanent mitigation")
-    else:
-        err(json.dumps(result, indent=4) + "\nunexpected result")
-elif action == "disable":
-    result = client.delete(f"/ip/{ip}/mitigation/{ip}")
-    if result["permanent"] is True and result["state"] == "removalPending":
-        print("disabling permanent mitigation")
-    else:
-        err(json.dumps(result, indent=4) + "\nunexpected result")
+        err("unknown action: " + action)
 else:
-    err("unknown action: " + action)
+        err("unknown system: " + system)