358 lines
15 KiB
Python
358 lines
15 KiB
Python
import logging
|
|
import os
|
|
import yaml
|
|
from PyQt5.QtWidgets import QDialog, QVBoxLayout, QMessageBox, QLabel, QRadioButton, QPushButton
|
|
from logline_leviathan.gui.ui_helper import UIHelper
|
|
|
|
from logline_leviathan.database.database_manager import *
|
|
|
|
|
|
|
|
|
|
class DatabaseOperations:
|
|
def __init__(self, main_window, db_init_func):
|
|
self.main_window = main_window
|
|
self.db_init_func = db_init_func
|
|
self.selected_resolutions = []
|
|
|
|
|
|
def ensureDatabaseExists(self):
|
|
db_path = 'entities.db'
|
|
db_exists = os.path.exists(db_path)
|
|
if not db_exists:
|
|
logging.info("Database does not exist. Creating new database...")
|
|
self.db_init_func() # This should call create_database
|
|
else:
|
|
logging.info("Database exists.")
|
|
|
|
def loadRegexFromYAML(self):
|
|
with open('./data/entities.yaml', 'r') as file:
|
|
yaml_data = yaml.safe_load(file)
|
|
clean_yaml_data = self.notify_duplicates_from_yaml(yaml_data)
|
|
return clean_yaml_data
|
|
|
|
def notify_duplicates_from_yaml(self, yaml_data):
|
|
duplicates = []
|
|
seen_fields = {'entity_type': {}, 'gui_name': {}, 'gui_tooltip': {}, 'regex_pattern': {}, 'script_parser': {}}
|
|
|
|
for entity_name, entity_data in yaml_data.items():
|
|
# Iterate through each field and check for duplicates
|
|
for field in seen_fields:
|
|
value = entity_data.get(field)
|
|
if value: # Only check non-empty values
|
|
if value in seen_fields[field]:
|
|
duplicates.append({
|
|
"duplicate_field": field,
|
|
"entity_name": entity_name,
|
|
"original_entity_name": seen_fields[field][value]
|
|
})
|
|
seen_fields[field][value] = entity_name
|
|
|
|
if duplicates:
|
|
self.show_duplicate_error_dialog(duplicates)
|
|
raise ValueError("Duplicate entries found in YAML file. Aborting.")
|
|
return yaml_data
|
|
|
|
def show_duplicate_error_dialog(self, duplicates):
|
|
dialog = DuplicateErrorDialog(duplicates)
|
|
dialog.exec_()
|
|
|
|
|
|
def show_resolve_inconsistencies_dialog(self, db_entity, yaml_entity):
|
|
dialog = ResolveInconsistenciesDialog([(db_entity, yaml_entity)])
|
|
result = dialog.exec_()
|
|
if result == QDialog.Accepted:
|
|
resolutions = dialog.getSelectedResolutions()
|
|
if resolutions:
|
|
return resolutions[0] # Return the first (and only) resolution
|
|
return None
|
|
|
|
|
|
def populate_and_update_entities_from_yaml(self, yaml_data):
|
|
with session_scope() as session:
|
|
db_entities = session.query(EntityTypesTable).all()
|
|
db_entity_dict = {entity.entity_type: entity for entity in db_entities}
|
|
|
|
for entity_name, entity_data in yaml_data.items():
|
|
entity_type = entity_data['entity_type']
|
|
db_entity = db_entity_dict.get(entity_type)
|
|
|
|
|
|
if db_entity is None:
|
|
db_entity = self.find_potentially_modified_entity(db_entities, entity_data)
|
|
|
|
if db_entity:
|
|
parser_enabled_db = db_entity.parser_enabled
|
|
entity_data['parser_enabled'] = parser_enabled_db
|
|
if self.is_duplicate_or_inconsistent(db_entity, entity_data, db_entities):
|
|
logging.warning(f"Issue found with entity {db_entity} and {entity_data}. Handling resolution.")
|
|
resolution = self.show_resolve_inconsistencies_dialog(db_entity, entity_data)
|
|
if resolution:
|
|
self.apply_resolution([(resolution, db_entity)], session) # Pass db_entity as part of the resolution
|
|
else:
|
|
for key, value in entity_data.items():
|
|
setattr(db_entity, key, value)
|
|
else:
|
|
new_entity = EntityTypesTable(**entity_data)
|
|
session.add(new_entity)
|
|
|
|
session.commit()
|
|
|
|
def find_potentially_modified_entity(self, db_entities, yaml_entity):
|
|
for db_ent in db_entities:
|
|
if any(
|
|
getattr(db_ent, key) == yaml_entity[key]
|
|
for key in ['entity_type', 'gui_name', 'gui_tooltip', 'regex_pattern', 'script_parser', 'parser_enabled']
|
|
if yaml_entity[key]
|
|
):
|
|
return db_ent
|
|
return None
|
|
|
|
|
|
|
|
def is_duplicate_or_inconsistent(self, db_entity, yaml_entity, db_entities):
|
|
if db_entity:
|
|
# Exclude 'parser_enabled' from the inconsistency check
|
|
keys_to_check = ['entity_type', 'gui_name', 'gui_tooltip', 'regex_pattern', 'script_parser']
|
|
for key in keys_to_check:
|
|
if getattr(db_entity, key, None) != yaml_entity.get(key) and yaml_entity.get(key) is not None:
|
|
logging.debug(f"Found inconsistent entity: DB-Entity: {db_entity} YAML-Entity: {yaml_entity}")
|
|
return True
|
|
|
|
# Check for duplicate across all entities
|
|
for db_ent in db_entities:
|
|
if db_ent.entity_type == yaml_entity['entity_type']:
|
|
continue
|
|
|
|
if any(
|
|
getattr(db_ent, key) == yaml_entity[key] and yaml_entity[key] is not None
|
|
for key in ['entity_type', 'gui_name', 'gui_tooltip', 'regex_pattern', 'script_parser',]
|
|
):
|
|
logging.debug(f"Found duplicate entity: {db_ent}")
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
def update_database_entry(self, db_entity, yaml_entity):
|
|
for key, value in yaml_entity.items():
|
|
setattr(db_entity, key, value)
|
|
|
|
|
|
|
|
def apply_resolution(self, resolutions, session):
|
|
with open('./data/entities.yaml', 'r') as file:
|
|
yaml_data = yaml.safe_load(file)
|
|
|
|
for (resolution, entity), db_entity in resolutions:
|
|
if resolution == 'yaml':
|
|
logging.debug(f"Resolving YAML entity: {entity} with resolution: yaml and db_entity: {db_entity}")
|
|
if db_entity:
|
|
foreign_keys = self.capture_foreign_keys(db_entity.entity_type_id, session)
|
|
session.delete(db_entity)
|
|
|
|
new_entity = EntityTypesTable(**entity)
|
|
session.add(new_entity)
|
|
session.flush()
|
|
|
|
self.reassign_foreign_keys(new_entity, foreign_keys, session)
|
|
|
|
elif resolution == 'db':
|
|
if entity: # Existing database entity is chosen
|
|
yaml_data[entity.entity_type] = {
|
|
'entity_type': entity.entity_type,
|
|
'gui_name': entity.gui_name,
|
|
'gui_tooltip': entity.gui_tooltip,
|
|
'parent_type': entity.parent_type,
|
|
'regex_pattern': entity.regex_pattern,
|
|
'script_parser': entity.script_parser,
|
|
'parser_enabled': entity.parser_enabled
|
|
}
|
|
|
|
with open('./data/entities.yaml', 'w') as file:
|
|
yaml.dump(yaml_data, file)
|
|
|
|
|
|
def capture_foreign_keys(self, entity_id, session):
|
|
foreign_keys = {}
|
|
|
|
# Use entity_id to capture references
|
|
distinct_entities_refs = session.query(DistinctEntitiesTable).filter_by(entity_types_id=entity_id).all()
|
|
foreign_keys['distinct_entities'] = [ref.distinct_entities_id for ref in distinct_entities_refs]
|
|
|
|
entities_refs = session.query(EntitiesTable).filter_by(entity_types_id=entity_id).all()
|
|
foreign_keys['entities'] = [ref.entities_id for ref in entities_refs]
|
|
|
|
return foreign_keys
|
|
|
|
|
|
def reassign_foreign_keys(self, new_entity, foreign_keys, session):
|
|
# Reassigning references in DistinctEntitiesTable
|
|
for distinct_id in foreign_keys.get('distinct_entities', []):
|
|
distinct_entity = session.query(DistinctEntitiesTable).get(distinct_id)
|
|
distinct_entity.entity_types_id = new_entity.entity_type_id
|
|
|
|
# Reassigning references in EntitiesTable
|
|
for entity_id in foreign_keys.get('entities', []):
|
|
entity = session.query(EntitiesTable).get(entity_id)
|
|
entity.entity_types_id = new_entity.entity_type_id
|
|
|
|
|
|
def checkScriptPresence(self):
|
|
parser_directory = './data/parser'
|
|
missing_scripts = []
|
|
|
|
with session_scope() as session:
|
|
all_entities = session.query(EntityTypesTable).all()
|
|
for entity in all_entities:
|
|
script_name = entity.script_parser
|
|
if script_name:
|
|
script_path = os.path.join(parser_directory, script_name)
|
|
if not os.path.exists(script_path):
|
|
missing_scripts.append(script_name)
|
|
|
|
if missing_scripts:
|
|
missing_scripts_str = "\n".join(missing_scripts)
|
|
msg = QMessageBox()
|
|
msg.setIcon(QMessageBox.Warning)
|
|
msg.setWindowTitle("Fehlende Skripte")
|
|
msg.setText(".\nDas ist nicht zwingend ein Problem, aber falls nötig,\nsollten die Skritpte in ./data/parser/ ergänzt werden.\nListe der erwarteten Skripte:")
|
|
msg.setInformativeText(missing_scripts_str)
|
|
msg.exec_() # Display the message box
|
|
|
|
return missing_scripts
|
|
|
|
def purgeWordlistEntries(self):
|
|
try:
|
|
with session_scope() as session:
|
|
# Identify the entity_type_id for 'generated_wordlist_match'
|
|
wordlist_entity_type = session.query(EntityTypesTable).filter_by(entity_type='generated_wordlist_match').one_or_none()
|
|
|
|
if not wordlist_entity_type:
|
|
logging.info("No 'generated_wordlist_match' entity type found. No action taken.")
|
|
return
|
|
|
|
# Find all distinct entities associated with the wordlist entity type
|
|
distinct_entities_to_remove = session.query(DistinctEntitiesTable).filter_by(entity_types_id=wordlist_entity_type.entity_type_id).all()
|
|
|
|
for distinct_entity in distinct_entities_to_remove:
|
|
# Remove all related entities entries and their context
|
|
entities_to_remove = session.query(EntitiesTable).filter_by(distinct_entities_id=distinct_entity.distinct_entities_id).all()
|
|
for entity in entities_to_remove:
|
|
# Remove related context entries
|
|
session.query(ContextTable).filter_by(entities_id=entity.entities_id).delete()
|
|
# Remove the entity itself
|
|
session.delete(entity)
|
|
|
|
# Commit the changes
|
|
session.commit()
|
|
except Exception as e:
|
|
logging.error(f"Error during wordlist entries purge: {str(e)}")
|
|
raise
|
|
|
|
|
|
class ResolveInconsistenciesDialog(QDialog):
|
|
def __init__(self, inconsistencies, parent=None):
|
|
super().__init__(parent)
|
|
self.setWindowTitle("Inkonsistenzen auflösen")
|
|
self.inconsistencies = inconsistencies
|
|
self.resolution_choices = []
|
|
self.selected_entity = None
|
|
self.selected_entities = []
|
|
self.selected_resolutions = []
|
|
self.initUI()
|
|
|
|
def initUI(self):
|
|
layout = QVBoxLayout(self)
|
|
|
|
for db_entity, yaml_entity in self.inconsistencies:
|
|
db_entity_str = self.format_entity_for_display(db_entity)
|
|
yaml_entity_str = self.format_entity_for_display(yaml_entity)
|
|
|
|
# Create labels and radio buttons for each inconsistency
|
|
db_label = QLabel(f"Datenbank-Eintrag: {db_entity_str}")
|
|
yaml_label = QLabel(f"YAML-Eintrag: {yaml_entity_str}")
|
|
db_radio = QRadioButton("Datenbank-Eintrag behalten")
|
|
yaml_radio = QRadioButton("YAML-Eintrag behalten")
|
|
|
|
layout.addWidget(db_label)
|
|
layout.addWidget(db_radio)
|
|
layout.addWidget(yaml_label)
|
|
layout.addWidget(yaml_radio)
|
|
|
|
self.resolution_choices.append((db_radio, yaml_radio))
|
|
|
|
# Buttons for OK and Cancel
|
|
btn_ok = QPushButton("OK", self)
|
|
btn_ok.clicked.connect(self.on_ok)
|
|
btn_cancel = QPushButton("Abbruch", self)
|
|
btn_cancel.clicked.connect(self.reject)
|
|
layout.addWidget(btn_ok)
|
|
layout.addWidget(btn_cancel)
|
|
|
|
def on_ok(self):
|
|
self.selected_resolutions = [] # Reset the list before storing new selections
|
|
for (db_radio, yaml_radio), (db_entity, yaml_entity) in zip(self.resolution_choices, self.inconsistencies):
|
|
if db_radio.isChecked():
|
|
self.selected_resolutions.append(('db', db_entity))
|
|
elif yaml_radio.isChecked():
|
|
self.selected_resolutions.append(('yaml', yaml_entity))
|
|
else:
|
|
self.selected_resolutions.append((None, None))
|
|
|
|
self.accept()
|
|
|
|
|
|
def getSelectedResolutions(self):
|
|
return self.selected_resolutions
|
|
|
|
def format_entity_for_display(self, entity):
|
|
if isinstance(entity, dict):
|
|
# YAML entity is already a dictionary
|
|
return "\n".join(f"{key}: {value}" for key, value in entity.items())
|
|
else:
|
|
# Database entity needs to be formatted
|
|
return "\n".join(f"{attr}: {getattr(entity, attr)}" for attr in ['entity_type', 'gui_name', 'gui_tooltip', 'parent_type', 'regex_pattern', 'script_parser', 'parser_enabled'])
|
|
|
|
|
|
|
|
class DuplicateErrorDialog(QDialog):
|
|
def __init__(self, duplicates, parent=None):
|
|
super().__init__(parent)
|
|
self.setWindowTitle("Duplikate gefunden")
|
|
self.duplicates = duplicates
|
|
self.initUI()
|
|
|
|
def initUI(self):
|
|
layout = QVBoxLayout(self)
|
|
|
|
# Display duplicate entries
|
|
error_label = QLabel("Duplikate wurden in ./data/entities.yaml gefunden. Diese sollten manuell aufgelöst werden:")
|
|
layout.addWidget(error_label)
|
|
|
|
for dup in self.duplicates:
|
|
dup_str = self.format_entity_for_display(dup)
|
|
dup_label = QLabel(dup_str)
|
|
layout.addWidget(dup_label)
|
|
|
|
# Buttons
|
|
open_button = QPushButton("YAML-Datei oeffnen", self)
|
|
open_button.clicked.connect(self.openYAML)
|
|
exit_button = QPushButton("Abbruch", self)
|
|
exit_button.clicked.connect(self.close)
|
|
|
|
layout.addWidget(open_button)
|
|
layout.addWidget(exit_button)
|
|
|
|
def format_entity_for_display(self, entity):
|
|
if isinstance(entity, dict):
|
|
return "\n".join(f"{key}: {value}" for key, value in entity.items())
|
|
|
|
def openYAML(self):
|
|
ui_helper = UIHelper(main_window=self)
|
|
ui_helper.openFile('data/entities.yaml')
|
|
|
|
|