class Database:
# _delays = [0.5, 0.5, 1, 1, 1]
_delays = [0.5]
_dbs = {}
_connected = False
class DatabaseNotReadyError(Exception):
pass
def __init__(self, url, username, password, dbs):
self._server = CouchDB(username,
password,
url=url,
auto_renew=True,
timeout=10)
try:
self._server.connect()
_connected = True
except ConnectionError:
logger.warning(
"Unable to connect database. Retry at next db call.")
for db in dbs:
if self._connected:
self._dbs[db] = self._server[db]
else:
self._dbs[db] = None
def __getattr__(self, item):
if item in self._dbs.keys():
if not self._connected:
try:
self._server.connect()
self._connected = True
except ConnectionError:
raise self.DatabaseNotReadyError("Database not connected.")
if self._dbs[item] is None:
self._dbs[item] = self._server[item]
return self._dbs[item]
else:
raise AttributeError
def update_doc(self, db, mango_query, data):
query = getattr(self, db).get_query_result(mango_query)
if query[0]:
doc = getattr(self, db)[query[0][0]["_id"]]
for item in data:
doc[item] = data[item]
try:
doc.save()
except HTTPError as e:
if e.response.status_code == 409:
doc.fetch()
for item in data:
doc[item] = data[item]
doc.save()
else:
data["_id"] = uuid4().hex
getattr(self, db).create_document(data)
class Database:
def __init__(self):
self.user = os.getenv("COUCHDB_USER")
self.password = os.getenv("COUCHDB_PASSWORD")
self.url = os.getenv("COUCHDB_URL")
self.connect = True
self.client = CouchDB(
**{
'user': self.user,
'auth_token': self.password,
'url': self.url,
'connect': self.connect
})
def get_users(self):
""" Returns a connection to the users database.
:return:
"""
users_db_name = 'users'
self.client.connect()
try:
users_db = self.client[users_db_name]
except KeyError as e:
print(e)
print(f"[CouchDB] Creating database: {users_db_name}")
users_db = self.client.create_database(users_db_name,
partitioned=False)
return users_db
def get_notes(self):
""" Returns a connection to the notes database.
Each note SHOULD BE partitioned to the appropriate user `_id` (aka email)
:return:
"""
notes_db_name = 'notes'
try:
notes_db = self.client[notes_db_name]
except KeyError as e:
print(e)
print(f"[CouchDB] Creating database: {notes_db_name}")
notes_db = self.client.create_database(notes_db_name,
partitioned=True)
return notes_db
from studium.core.models import User, NewUser, UserInDB, UserLogin
from studium.core.security import authenticate_user, get_current_active_user, get_password_hash
from studium.core.jwt import create_access_token, ACCESS_TOKEN_EXPIRE_MINUTES
sentry_dsn = os.getenv("SENTRY_DSN")
sentry_sdk.init(sentry_dsn, traces_sample_rate=1.0)
client = CouchDB(
**{
'user': os.getenv("COUCHDB_USER"),
'auth_token': os.getenv("COUCHDB_PASSWORD"),
'url': os.getenv("COUCHDB_URL"),
'connect': True
})
client.connect()
# Database connection
db = Database()
users_db = db.get_users()
app = FastAPI(debug=True)
origins = [
"http://localhost.tiangolo.com",
"https://localhost.tiangolo.com",
"http://localhost",
"http://localhost:8080",
]
class ResultDatabase():
def __init__(self, couchdb_host, couchdb_database, couchdb_username,
couchdb_password):
self.couchdb = CouchDB(couchdb_username,
couchdb_password,
url=couchdb_host)
self.couchdb_database = couchdb_database
def _get_db_handle(self):
self.couchdb.connect()
return self.couchdb[self.couchdb_database]
def _get_ddoc(self):
couchdb_database = self._get_db_handle()
return couchdb_database.get_design_document("results")
def _get_result_history(self):
result_ddoc = self._get_ddoc()
return result_ddoc.get_view("check_result_history")
def _get_results(self):
result_ddoc = self._get_ddoc()
return result_ddoc.get_view("check_results")
def _get_host_status_history(self):
result_ddoc = self._get_ddoc()
return result_ddoc.get_view("host_status_history")
def _get_host_status(self):
result_ddoc = self._get_ddoc()
return result_ddoc.get_view("host_status")
def _get_severity(self):
result_ddoc = self._get_ddoc()
return result_ddoc.get_view("severity")
def _get_scenario_severity(self):
result_ddoc = self._get_ddoc()
return result_ddoc.get_view("scenario_severity")
def _get_successful_checks(self):
result_ddoc = self._get_ddoc()
return result_ddoc.get_view("successful_checks")
def host_status(self, username=None, hostname=None):
host_status_view = self._get_host_status()
if username is None and hostname is None:
host_status = host_status_view()
elif hostname is None:
host_status = host_status_view(startkey=[username],
endkey=[username, {}])
else:
host_status = host_status_view(keys=[[username, hostname]])
return host_status["rows"]
def check_results(self, username=None, check_name=None):
check_results_view = self._get_results()
if username is None and check_name is None:
check_results = check_results_view()
elif check_name is None:
check_results = check_results_view(startkey=[username],
endkey=[username, {}])
else:
if isinstance(check_name, list):
keys = map(lambda x: [username, x], check_name)
else:
keys = [[username, check_name]]
check_results = check_results_view(keys=keys)
return check_results["rows"]
def active_students(self, scenario_name):
successful_checks_view = self._get_successful_checks()
students = map(
lambda x: x["key"][1],
successful_checks_view(startkey=[scenario_name],
endkey=[scenario_name, {}],
group_level=2,
reduce=True)["rows"])
return students
def successful_checks_count(self, scenario_name, username):
return len(self.successful_checks(scenario_name, username))
def successful_checks(self, scenario_name, username):
successful_checks_view = self._get_successful_checks()
successful_checks_rows = successful_checks_view(
startkey=[scenario_name, username],
endkey=[scenario_name, username, {}],
group_level=3,
reduce=True)["rows"]
return map(lambda x: [x["key"][2], x["value"]], successful_checks_rows)
def get_attachment(self, username, checkname, attachment_name):
document_id = "attachments_" + username + "_" + checkname
couchdb_database = self._get_db_handle()
if document_id in couchdb_database:
document = couchdb_database[document_id]
return document.get_attachment(attachment_name)
else:
return dict()