From 44625db6973c110c38d7129966e4e09c10945ec4 Mon Sep 17 00:00:00 2001
From: Stefan Midjich <swehack@gmail.com>
Date: Fri, 29 Sep 2017 17:44:40 +0200
Subject: [PATCH] added "expires" field to clients

---
 tools/captiveportal.pgsql |  1 +
 tools/client.py           |  5 +++--
 tools/manage_client.py    | 28 ++++++++++++++++++++++++++++
 tools/storage.py          | 11 ++++++-----
 4 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/tools/captiveportal.pgsql b/tools/captiveportal.pgsql
index e7a4555..7043551 100644
--- a/tools/captiveportal.pgsql
+++ b/tools/captiveportal.pgsql
@@ -16,6 +16,7 @@ create table if not exists client (
  enabled boolean NOT NULL,
  last_packets bigint default 0,
  last_activity timestamp,
+ expires timestamp DEFAULT NULL,
  primary key (client_id, ip_address, protocol)
 );
 
diff --git a/tools/client.py b/tools/client.py
index 42e1486..6121d01 100644
--- a/tools/client.py
+++ b/tools/client.py
@@ -4,11 +4,11 @@ Handles "clients" in IPtables for captive portal.
 
 import ipaddress
 from uuid import uuid4
-from datetime import datetime
+from datetime import datetime, timedelta
 
 import iptc
 
-import errors
+from errors import StorageNotFound
 
 
 class Client(object):
@@ -46,6 +46,7 @@ class Client(object):
             self.enabled = False
             self.last_packets = 0
             self.last_activity = None
+            self.expires = datetime.now() + timedelta(days=1)
 
         # Init iptables
         self.table = iptc.Table(iptc.Table.MANGLE)
diff --git a/tools/manage_client.py b/tools/manage_client.py
index fd1d646..c0782d2 100644
--- a/tools/manage_client.py
+++ b/tools/manage_client.py
@@ -6,17 +6,45 @@ from sys import exit
 from argparse import ArgumentParser, FileType
 from pprint import pprint as pp
 from configparser import RawConfigParser
+from datetime import datetime
 
 import errors
 from storage import StoragePostgres
 from client import Client
 
+
+# Custom defined argparse types for dates
+def valid_date_type(arg_date_str):
+    """custom argparse *date* type for user dates values given from the command line"""
+    try:
+        return datetime.datetime.strptime(arg_date_str, "%Y-%m-%d")
+    except ValueError:
+        msg = "Given Date ({0}) not valid! Expected format, YYYY-MM-DD!".format(arg_date_str)
+        raise argparse.ArgumentTypeError(msg)
+
+
+def valid_datetime_type(arg_datetime_str):
+    """custom argparse type for user datetime values given from the command line"""
+    try:
+        return datetime.datetime.strptime(arg_datetime_str, "%Y-%m-%d %H:%M")
+    except ValueError:
+        msg = "Given Datetime ({0}) not valid! Expected format, 'YYYY-MM-DD HH:mm'!".format(arg_datetime_str)
+        raise argparse.ArgumentTypeError(msg) 
+
+
 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(
+    '--expires',
+    type=valid_datetime_type,
+    default=datetime.now() + timedelta(days=1),
+    help='Expiry date in format "YYYY-MM-DD HH:mm"'
+)
+
 parser.add_argument(
     '--disable',
     default=False,
diff --git a/tools/storage.py b/tools/storage.py
index ddd3036..0da4e53 100644
--- a/tools/storage.py
+++ b/tools/storage.py
@@ -56,10 +56,10 @@ class StoragePostgres(object):
     def write_client(self, client):
         query = (
             'insert into client (client_id, created, ip_address, protocol, '
-            'enabled, last_packets, last_activity) values (%s, %s, %s, %s, '
-            '%s, %s, %s) on conflict (client_id, ip_address, protocol) do '
-            'update set (enabled, last_packets, last_activity) = '
-            '(EXCLUDED.enabled, EXCLUDED.last_packets, '
+            'enabled, last_packets, last_activity, expires) values '
+            '(%s, %s, %s, %s, %s, %s, %s, %s) on conflict (client_id, '
+            'ip_address, protocol) do update set (enabled, last_packets, '
+            'last_activity) = (EXCLUDED.enabled, EXCLUDED.last_packets, '
             'EXCLUDED.last_activity)'
         )
         self.cur.execute(
@@ -71,7 +71,8 @@ class StoragePostgres(object):
                 client.protocol,
                 client.enabled,
                 client.last_packets,
-                client.last_activity
+                client.last_activity,
+                client.expires
             )
         )
         self.conn.commit()