blocksoft/app/libraries/database.py
2023-04-13 18:10:58 +02:00

286 líneas
7,1 KiB
Python

import os
import sys
import psycopg2
from psycopg2 import sql
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
import uuid
from datetime import datetime
import pytz
import pdb
tz = pytz.timezone('Europe/Madrid')
class Database():
name = 'fediverse database library'
def __init__(self, config_file=None, blocker_db=None, blocker_db_user=None, blocker_db_user_password=None):
self.config_file = "config/db_config.txt"
self.blocker_db = self.__get_parameter("blocker_db", self.config_file)
self.blocker_db_user = self.__get_parameter("blocker_db_user", self.config_file)
self.blocker_db_user_password = self.__get_parameter("blocker_db_user_password", self.config_file)
db_setup = self.__check_dbsetup(self)
if not db_setup:
self.blocker_db = input("\nblocksoft database name: ")
self.blocker_db_user = input("\nblocksoft database user: ")
self.blocker_db_user_password = input("\nblocksoft database user password: ")
self.__createdb(self)
self.__create_config(self)
self.__write_config(self)
def get_servers(self, software):
servers_list = []
try:
conn = None
conn = psycopg2.connect(database = self.blocker_db, user = self.blocker_db_user, password = "", host = "/var/run/postgresql", port = "5432")
cur = conn.cursor()
cur.execute("select server from servers where software=(%s)", (software,))
rows = cur.fetchall()
for row in rows:
servers_list.append(row[0])
cur.close()
except (Exception, psycopg2.DatabaseError) as error:
sys.exit(error)
finally:
if conn is not None:
conn.close()
return (servers_list)
def write_server(self, server, api, software):
now = datetime.now()
insert_sql = "INSERT INTO servers(server, api, software, created_at) VALUES(%s, %s, %s, %s) ON CONFLICT DO NOTHING"
conn = None
try:
conn = psycopg2.connect(database=self.blocker_db, user=self.blocker_db_user, password = self.blocker_db_user_password, host="/var/run/postgresql", port="6432")
cur = conn.cursor()
cur.execute(insert_sql, (server, api, software, now))
conn.commit()
cur.close()
except (Exception, psycopg2.DatabaseError) as error:
print(error)
finally:
if conn is not None:
conn.close()
def update_peer(self, server):
now = datetime.now()
update_sql = "UPDATE servers set updated_at=(%s) where server=(%s)"
conn = None
try:
conn = psycopg2.connect(database=self.blocker_db, user=self.blocker_db_user, password = self.blocker_db_user_password, host="/var/run/postgresql", port="6432")
cur = conn.cursor()
cur.execute(update_sql, (now, server))
conn.commit()
cur.close()
except (Exception, psycopg2.DatabaseError) as error:
print(error)
finally:
if conn is not None:
conn.close()
def update_server(self, server, api):
now = datetime.now()
update_sql = "UPDATE servers set server=(%s), api=(%s), updated_at=(%s) where server=(%s)"
conn = None
try:
conn = psycopg2.connect(database=self.blocker_db, user=self.blocker_db_user, password = self.blocker_db_user_password, host="/var/run/postgresql", port="6432")
cur = conn.cursor()
cur.execute(update_sql, (server, api, now, server))
conn.commit()
cur.close()
except (Exception, psycopg2.DatabaseError) as error:
print(error)
finally:
if conn is not None:
conn.close()
@staticmethod
def __check_dbsetup(self):
db_setup = False
try:
conn = None
conn = psycopg2.connect(database = self.blocker_db, user = self.blocker_db_user, password = self.blocker_db_user_password, host = "/var/run/postgresql", port = "6432")
db_setup = True
except (Exception, psycopg2.DatabaseError) as error:
print(error)
return db_setup
@staticmethod
def __createdb(self):
conn = None
try:
conn = psycopg2.connect(dbname='postgres',
user=self.blocker_db_user, host='',
password=self.blocker_db_user_password)
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
cur = conn.cursor()
print(f"Creating database {self.blocker_db}. Please wait...")
cur.execute(sql.SQL("CREATE DATABASE {}").format(
sql.Identifier(self.blocker_db))
)
print(f"Database {self.blocker_db} created!\n")
self.__dbtables_schemes(self)
except (Exception, psycopg2.DatabaseError) as error:
print(error)
finally:
if conn is not None:
conn.close()
@staticmethod
def __dbtables_schemes(self):
db = self.blocker_db
table = "servers"
sql = "create table "+table+" (server varchar(200) PRIMARY KEY, api varchar(50), software varchar(50), created_at timestamptz, updated_at timestamptz, downs int)"
self.__create_table(self, table, sql)
@staticmethod
def __create_table(self, table, sql):
conn = None
try:
conn = psycopg2.connect(database = self.blocker_db, user = self.blocker_db_user, password = self.blocker_db_user_password, host = "/var/run/postgresql", port = "6432")
cur = conn.cursor()
print(f"Creating table {table}")
cur.execute(sql)
conn.commit()
print(f"Table {table} created!\n")
except (Exception, psycopg2.DatabaseError) as error:
print(error)
finally:
if conn is not None:
conn.close()
def __get_parameter(self, parameter, config_file):
if not os.path.isfile(config_file):
print(f"File {config_file} not found..")
return
with open( config_file ) as f:
for line in f:
if line.startswith( parameter ):
return line.replace(parameter + ":", "").strip()
print(f"{config_file} Missing parameter {parameter}")
sys.exit(0)
@staticmethod
def __create_config(self):
if not os.path.exists('config'):
os.makedirs('config')
if not os.path.exists(self.config_file):
print(self.config_file + " created!")
with open(self.config_file, 'w'): pass
@staticmethod
def __write_config(self):
with open(self.config_file, 'a') as the_file:
the_file.write(f'blocker_db: {self.blocker_db}\nblocker_db_user: {self.blocker_db_user}\nblocker_db_user_password: {self.blocker_db_user_password}')
print(f"adding parameters to {self.config_file}\n")