XylotrechusZ
import logging
import itertools
import os
from contextlib import suppress
from datetime import datetime
from shutil import copy
from sqlite3 import connect, DatabaseError
from playhouse.sqlite_ext import SqliteExtDatabase
from defence360agent.application import app
from defence360agent import simple_rpc
from defence360agent.contracts.config import Model
from defence360agent.model import simplification
logger = logging.getLogger(__name__)
class OperationError(Exception):
pass
WORKAROUND_MSG = "Blank database will be created on agent start "
def check_and_repair():
base = Model.PATH
if simple_rpc.is_running():
raise OperationError(
"Cannot perform database check and backup while agent is running. "
"Please, stop the imunify360 agent with `service imunify360 stop`"
)
elif not os.path.isfile(base):
raise OperationError(
"DB %s is not exists. %s" % (base, WORKAROUND_MSG)
)
else:
if is_db_corrupted(base):
backup = make_backup(base)
if not backup:
raise OperationError(
"Cannot proceed without backup copy of the database."
"Please contact imunify360 support team at "
"https://cloudlinux.zendesk.com"
)
dump = dump_to_sql(base)
logger.info("Removing original corrupted database at %s" % base)
# TODO: Notify user in UI that original DB was dropped
os.remove(base)
if not dump:
raise OperationError(
"Cannot dump database to sql. Old DB backuped at %s. %s"
% (backup, WORKAROUND_MSG)
)
else:
restored = load_from_sql(base, dump)
if not restored:
raise OperationError(
"Loading dump to new database failed. Database will "
"be recreated during migrations."
)
if is_db_corrupted(restored):
os.remove(restored)
raise OperationError(
"Restored database is still corrupt. Removing "
"restored database. %s" % WORKAROUND_MSG
)
logger.info(
"Database restored successfully. Removing dump %s" % dump
)
os.remove(dump)
try:
logger.info("Performing migrations on restored database")
simplification.migrate()
except Exception as e:
os.remove(base)
raise OperationError(
"Migrations on restored database failed: %s. %s"
% (e, WORKAROUND_MSG)
)
else:
if not all_tables_are_present():
os.remove(base)
raise OperationError(
"Restored database does not "
"contain all necessary tables. "
"%s" % WORKAROUND_MSG
)
def mark_with_timestamp(filename, extension=None):
"""
>>> mark_with_timestamp('/var/imunify360/imunify360.db')
'/var/imunify360/imunify360.db_2017-09-26_03:33:44.705967'
>>> mark_with_timestamp('/var/imunify360/imunify360.db', extension='sql')
'/var/imunify360/imunify360.db_2017-09-26_03:34:01.098544.sql'
"""
instant = datetime.now()
basename = "{}_{}".format(filename, instant.isoformat("_"))
if extension:
return basename + ".%s" % extension
else:
return basename
def is_db_corrupted(db_path):
logger.info("Database %s integrity check..." % db_path)
is_corrupted = True
with connect(db_path) as connection:
try:
cursor = connection.execute("PRAGMA INTEGRITY_CHECK;")
result = next(cursor)
if "ok" in result:
logger.info("Database integrity check succeeded.")
is_corrupted = False
except DatabaseError as e:
logger.warning("DatabaseError detected: %s", e)
return is_corrupted
def dump_to_sql(db_path):
dumpfile = mark_with_timestamp(db_path, extension="sql")
logger.info("Dumping imunify360 database to %s" % dumpfile)
try:
with open(dumpfile, "w") as dump, connect(db_path) as connection:
for row in connection.iterdump():
dump.write(row)
except (DatabaseError, OSError) as e:
logger.error("Error during dump: %s. Operation aborted" % e)
with suppress(OSError):
os.remove(dumpfile)
dumpfile = None
return dumpfile
def load_from_sql(db_path, dumpfile):
# This is unlikely to happen because we delete the original file
# but to be defensive here won't hurt in case of reuse in other places.
if os.path.exists(db_path):
logger.warning(
"Database already exists. Loading dump to existing "
"database may cause errors. Operation aborted"
)
return None
logger.info(
"Reading dump %s into new database %s..." % (dumpfile, db_path)
)
with open(dumpfile, "r") as dump, connect(db_path) as connection:
# We cannot read line by line because SQL statements are dumped
# not in a statement-per-line way
try:
sql = dump.read()
connection.executescript(sql)
except MemoryError as e:
logger.error(e)
with suppress(OSError):
os.remove(db_path)
db_path = None
return db_path
def make_backup(db_path):
logger.info("Making backup of the %s..." % db_path)
backup_filename = mark_with_timestamp(db_path, "backup")
try:
copy(db_path, backup_filename)
logger.info("Database copied successfully to: %s " % backup_filename)
except Exception as e:
logger.error("Making backup failed: %s", e)
with suppress(OSError):
os.remove(backup_filename)
backup_filename = None
return backup_filename
def all_tables_are_present():
logger.info(
"Verifying that db schema is up-to-date and all tables are present..."
)
models = itertools.chain(
*[
simplification.get_models(module)
for module in app.MODULES_WITH_MODELS
]
)
if all(model.table_exists() for model in models):
logger.info("All tables are present")
return True
else:
logger.error("Some tables are missing in db.")
return False
def recreate_schema() -> None:
simplification.instance.db.init(Model.PATH)
logger.info("Recreating schema for linked DBs...")
attached_schemas = []
for db_path, schema in app.MIGRATIONS_ATTACHED_DBS:
logger.info("Attach db: %s", db_path)
simplification.instance.db.execute_sql(
"ATTACH ? AS ?", (db_path, schema)
)
attached_schemas.append(schema)
recreate_schema_models(simplification.instance.db, attached_schemas)
logger.info("Schema recreated successfully.")
def recreate_schema_models(
db: SqliteExtDatabase, target_schemas: list[str]
) -> None:
models_to_create = [
model
for model in itertools.chain(
*[
simplification.get_models(module)
for module in app.MODULES_WITH_MODELS
]
)
if model._meta.schema in target_schemas
]
logger.info("%r", models_to_create)
# bind models to the db to avoid issues related to initialization order
db.bind(models_to_create)
db.create_tables(models_to_create)
logger.info("Schema models recreated successfully.")