""" Database storage backends for client.py. """ import json from datetime import datetime import psycopg2 from psycopg2.extras import DictCursor, register_ipaddress, Inet from redis import Redis from portalclientlib.client import Client class StoragePostgres(object): """ This requires python 3 for inet data type. """ 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'), port=config.getint('postgres', 'port'), sslmode='disable', cursor_factory=DictCursor ) self.cur = self.conn.cursor() register_ipaddress() def client_ids(self): self.cur.execute( 'select client_id from authenticated_clients' ) return self.cur.fetchall() def get_client_by_id(self, client_id): self.cur.execute( 'select * from authenticated_clients where client_id=%s', (client_id,) ) return self.cur.fetchone() def get_client(self, ip_address): """ Expects an ipaddress.IPv4Interface as ip_address argument. """ self.cur.execute( 'select * from authenticated_clients where ip_address=%s', (ip_address, ) ) return self.cur.fetchone() def write_client(self, client): query = ( 'insert into authenticated_clients (client_id, created, ip_address, protocol, ' '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, expires) = (EXCLUDED.enabled, EXCLUDED.last_packets, ' 'EXCLUDED.last_activity, EXCLUDED.expires)' ) self.cur.execute( query, ( client.client_id, client.created, client.ip_address, client.protocol, client.enabled, client.last_packets, client.last_activity, client.expires ) ) self.conn.commit() def remove_client(self, client): query = 'delete from authenticated_clients where client_id=%s' self.cur.execute(query, (client.client_id,)) self.conn.commit() 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): """ Note: Abandoned this storage backend for Postgres. """ 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): raise NotImplemented