captive.whump.shanti-portal/tools/manage_client.py

111 lines
2.8 KiB
Python
Raw Normal View History

#!/usr/bin/env python
2017-02-22 13:41:58 +00:00
# Python helper tool to add IPtables rule using the iptc library. This must
# of course run as root for iptc to work.
2017-03-06 15:03:57 +00:00
from sys import exit
2017-09-29 16:40:57 +00:00
from argparse import ArgumentParser, FileType, ArgumentTypeError
from pprint import pprint as pp
2017-03-03 00:04:12 +00:00
from configparser import RawConfigParser
2017-09-29 15:56:16 +00:00
from datetime import datetime, timedelta
2017-03-06 15:03:57 +00:00
import errors
from storage import StoragePostgres
2017-03-03 00:04:12 +00:00
from client import Client
2017-09-29 15:44:40 +00:00
# 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:
2017-09-29 16:40:57 +00:00
return datetime.strptime(arg_date_str, "%Y-%m-%d")
2017-09-29 15:44:40 +00:00
except ValueError:
msg = "Given Date ({0}) not valid! Expected format, YYYY-MM-DD!".format(arg_date_str)
2017-09-29 16:40:57 +00:00
raise ArgumentTypeError(msg)
2017-09-29 15:44:40 +00:00
def valid_datetime_type(arg_datetime_str):
"""custom argparse type for user datetime values given from the command line"""
try:
2017-09-29 16:40:57 +00:00
return datetime.strptime(arg_datetime_str, "%Y-%m-%d %H:%M")
2017-09-29 15:44:40 +00:00
except ValueError:
msg = "Given Datetime ({0}) not valid! Expected format, 'YYYY-MM-DD HH:mm'!".format(arg_datetime_str)
2017-09-29 16:40:57 +00:00
raise ArgumentTypeError(msg)
2017-09-29 15:44:40 +00:00
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.'
))
2017-09-29 15:44:40 +00:00
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,
2017-03-07 16:21:52 +00:00
action='store_true',
help='Disable the client in the DB and delete from firewall'
)
parser.add_argument(
'--delete',
default=False,
2017-03-07 16:21:52 +00:00
action='store_true',
help='Delete the client from DB and firewall'
)
2017-03-03 00:04:12 +00:00
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()
2017-03-03 00:04:12 +00:00
config = RawConfigParser()
config.readfp(args.config)
2017-03-06 15:03:57 +00:00
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('Client not found')
2017-03-06 15:03:57 +00:00
exit(1)
if args.disable:
2017-09-29 16:50:41 +00:00
enabled = False
else:
enabled = True
if args.delete:
# This both deletes the iptables rule AND the client entry from DB.
client.delete()
else:
2017-09-29 16:50:41 +00:00
if args.expires:
client.expires = args.expires
client.enabled = enabled
client.commit()