From 3059459915081a2f67d1e28d5403cb72dd28d5a8 Mon Sep 17 00:00:00 2001
From: Stefan Midjich <swehack@gmail.com>
Date: Tue, 7 Mar 2017 17:15:39 +0100
Subject: [PATCH] merged three actions into manage_client.py

---
 tools/client.py                           |  4 +-
 tools/{add_client.py => manage_client.py} | 31 ++++++++++++--
 tools/remove_client.py                    | 50 -----------------------
 3 files changed, 30 insertions(+), 55 deletions(-)
 rename tools/{add_client.py => manage_client.py} (57%)
 delete mode 100644 tools/remove_client.py

diff --git a/tools/client.py b/tools/client.py
index 210534e..1636652 100644
--- a/tools/client.py
+++ b/tools/client.py
@@ -63,7 +63,9 @@ class Client(object):
 
     def commit(self):
         self.commit_client()
-        self.commit_rule()
+
+        if self.enabled:
+            self.commit_rule()
 
 
     def commit_client(self):
diff --git a/tools/add_client.py b/tools/manage_client.py
similarity index 57%
rename from tools/add_client.py
rename to tools/manage_client.py
index c4f4ca5..f5f1954 100644
--- a/tools/add_client.py
+++ b/tools/manage_client.py
@@ -11,7 +11,23 @@ import errors
 from storage import StoragePostgres
 from client import Client
 
-parser = ArgumentParser()
+parser = ArgumentParser((
+    'Handle clients in the captive portal. Default mode of operation is to'
+    ' create new clients and enable them. Other mode is to --disable the '
+    'client. And last mode is to --delete the client completely.'
+))
+
+parser.add_argument(
+    '--disable',
+    default=False,
+    help='Disable the client in the DB and delete from firewall'
+)
+
+parser.add_argument(
+    '--delete',
+    default=False,
+    help='Delete the client from DB and firewall'
+)
 
 parser.add_argument(
     '--protocol',
@@ -46,8 +62,15 @@ try:
         chain=config.get('iptables', 'chain')
     )
 except errors.StorageNotFound:
-    print('Could not find client')
+    print('Client not found')
     exit(1)
 
-client.enabled = True
-client.commit()
+if args.disable:
+    # For non-existing clients this actually creates them in disabled mode.
+    client.enabled = False
+    client.commit()
+elif args.delete:
+    client.delete()
+else:
+    client.enabled = True
+    client.commit()
diff --git a/tools/remove_client.py b/tools/remove_client.py
deleted file mode 100644
index 9f1e815..0000000
--- a/tools/remove_client.py
+++ /dev/null
@@ -1,50 +0,0 @@
-#!/usr/bin/env python
-
-from sys import exit
-from argparse import ArgumentParser, FileType
-from pprint import pprint as pp
-from configparser import RawConfigParser
-
-import errors
-from storage import StoragePostgres
-from client import Client
-
-parser = ArgumentParser()
-
-parser.add_argument(
-    '--protocol',
-    required=True,
-    choices=['tcp', 'udp'],
-    help='Protocol for client'
-)
-
-parser.add_argument(
-    '--config',
-    type=FileType('r'),
-    required=True,
-    help='Configuration file'
-)
-
-parser.add_argument(
-    'src_ip',
-    help='Client source IP to add'
-)
-
-args = parser.parse_args()
-
-config = RawConfigParser()
-config.readfp(args.config)
-
-sr = StoragePostgres(config=config)
-try:
-    client = Client(
-        storage=sr,
-        ip_address=args.src_ip,
-        protocol=args.protocol,
-        chain=config.get('iptables', 'chain')
-    )
-except errors.StorageNotFound:
-    print('Could not find client')
-    exit(1)
-
-client.delete()