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

110 lines
2.9 KiB
Python
Raw Normal View History

2017-03-03 00:04:12 +00:00
"""
Database storage backends for client.py.
"""
import json
from datetime import datetime
2017-03-06 09:26:35 +00:00
import psycopg2
2017-03-06 15:03:57 +00:00
from psycopg2.extras import DictCursor, register_ipaddress, Inet
2017-03-03 00:04:12 +00:00
from redis import Redis
2017-03-06 15:03:57 +00:00
from client import Client
2017-03-03 00:04:12 +00:00
2017-03-06 09:26:35 +00:00
class StoragePostgres(object):
2017-03-06 15:03:57 +00:00
"""
This requires python 3 for inet data type.
"""
2017-03-06 09:26:35 +00:00
def __init__(self, **kw):
config = kw.pop('config')
self.conn = psycopg2.connect(
host=config.get('postgres', 'hostname'),
user=config.get('postgres', 'username'),
password=config.get('postgres', 'password'),
dbname=config.get('postgres', 'database'),
2017-03-06 15:03:57 +00:00
port=config.getint('postgres', 'port'),
sslmode='disable',
cursor_factory=DictCursor
)
self.cur = self.conn.cursor()
register_ipaddress()
def get_client_by_id(self, client_id):
self.cur.execute(
'select * from client where client_id=%s',
(client_id,)
2017-03-06 09:26:35 +00:00
)
2017-03-06 15:03:57 +00:00
return self.cur.fetchone()
def get_client(self, ip_address, protocol):
self.cur.execute(
'select * from client where ip_address=%s and protocol=%s',
(Inet(ip_address), protocol, )
)
return self.cur.fetchone()
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, '
'EXCLUDED.last_activity)'
)
self.cur.execute(
query,
(
client.client_id,
client.created,
client.ip_address,
client.protocol,
client.enabled,
client.last_packets,
client.last_activity
)
)
self.conn.commit()
def remove_client(self, client):
query = 'delete from client where client_id=%s'
self.cur.execute(query, (client.client_id,))
self.conn.commit()
2017-03-06 09:26:35 +00:00
2017-03-03 00:04:12 +00:00
class DateTimeEncoder(json.JSONEncoder):
"""
json.JSONEncoder sub-class that converts all datetime objects to
epoch timestamp integer values.
"""
def default(self, o):
if isinstance(o, datetime):
return int(o.strftime('%s'))
return json.JSONEncoder.default(self, o)
class StorageRedis(object):
2017-03-06 09:26:35 +00:00
"""
Note: Abandoned this storage backend for Postgres.
"""
2017-03-03 00:04:12 +00:00
def __init__(self, **kw):
config = kw.pop('config')
self.r = Redis(
host=config.get('redis', 'hostname'),
port=config.getint('redis', 'port'),
db=config.getint('redis', 'db')
)
def add_client(self, client_id, **kw):
2017-03-06 09:26:35 +00:00
raise NotImplemented