From 49f34b8124fc0344f09fbfaf2d3060a9464e18b5 Mon Sep 17 00:00:00 2001 From: ibuler Date: Wed, 8 Jan 2025 19:13:13 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E9=87=8D=E6=9E=84=E6=89=AB=E6=8F=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../accounts/api/automations/check_account.py | 29 +- .../check_account/leak_passwords.db | 2 +- .../automations/check_account/manager.py | 398 +++++++++--------- apps/assets/automations/base/manager.py | 16 +- 4 files changed, 229 insertions(+), 216 deletions(-) diff --git a/apps/accounts/api/automations/check_account.py b/apps/accounts/api/automations/check_account.py index cbc7a524d..34e921c62 100644 --- a/apps/accounts/api/automations/check_account.py +++ b/apps/accounts/api/automations/check_account.py @@ -2,12 +2,10 @@ # from django.db.models import Q, Count from django.http import HttpResponse -from rest_framework.decorators import action -from rest_framework.exceptions import MethodNotAllowed -from operator import itemgetter from django.shortcuts import get_object_or_404 from django.utils import timezone - +from rest_framework.decorators import action +from rest_framework.exceptions import MethodNotAllowed from rest_framework.response import Response from accounts import serializers @@ -140,7 +138,7 @@ class CheckAccountEngineViewSet(JMSModelViewSet): serializer_class = serializers.CheckAccountEngineSerializer @staticmethod - def init_if_need(): + def get_default_engines(): data = [ { "id": "00000000-0000-0000-0000-000000000001", @@ -154,15 +152,32 @@ class CheckAccountEngineViewSet(JMSModelViewSet): "name": "检查账号密码强弱", "comment": "基于账号密码的安全性进行检查分析, 检查密码强度、泄露等信息", }, + { + "id": "00000000-0000-0000-0000-000000000003", + "slug": "check_account_repeat", + "name": "检查账号密码是否重复", + "comment": "检查账号是否与其它账号相同" + }, + { + "id": "00000000-0000-0000-0000-000000000004", + "slug": "check_account_leak", + "name": "检查账号密码是否是常见密码", + "comment": "检查账号密码是否是常见泄露的密码" + }, ] + return data + def init_if_need(self): + data = self.get_default_engines() model_cls = CheckAccountEngine - if model_cls.objects.all().count() == 2: + + if model_cls.objects.count() == 4: return for item in data: - model_cls.objects.create(**item) + model_cls.objects.update_or_create(defaults=item, id=item["id"]) def get_queryset(self): + # return self.get_default_engines() self.init_if_need() return CheckAccountEngine.objects.all() diff --git a/apps/accounts/automations/check_account/leak_passwords.db b/apps/accounts/automations/check_account/leak_passwords.db index c5ac86a2f..e241a81e3 100644 --- a/apps/accounts/automations/check_account/leak_passwords.db +++ b/apps/accounts/automations/check_account/leak_passwords.db @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:8526225e13163cc92535a0ed3ef1c12cfde471173394504aea2b0583fe68efb5 +oid sha256:37be0604185cf3189de457dc7339a461052d96582096beb4fdf7b0fcb2c290a8 size 98037760 diff --git a/apps/accounts/automations/check_account/manager.py b/apps/accounts/automations/check_account/manager.py index 2575fcb10..294303e70 100644 --- a/apps/accounts/automations/check_account/manager.py +++ b/apps/accounts/automations/check_account/manager.py @@ -1,42 +1,16 @@ +import hashlib import os import re import sqlite3 +import uuid from django.conf import settings from django.utils import timezone from accounts.models import Account, AccountRisk, RiskChoice from assets.automations.base.manager import BaseManager +from common.const import ConfirmOrIgnore from common.decorators import bulk_create_decorator, bulk_update_decorator -from common.utils.strings import color_fmt - - -def is_weak_password(password): - # 判断密码长度 - if len(password) < 8: - return True - - # 判断是否只有一种字符类型 - if password.isdigit() or password.isalpha(): - return True - - # 判断是否只包含数字或字母 - if password.islower() or password.isupper(): - return True - - # 判断是否包含常见弱密码 - common_passwords = ["123456", "password", "12345678", "qwerty", "abc123"] - if password.lower() in common_passwords: - return True - - # 正则表达式判断字符多样性（数字、字母、特殊字符） - if ( - not re.search(r"[A-Za-z]", password) - or not re.search(r"[0-9]", password) - or not re.search(r"[\W_]", password) - ): - return True - return False @bulk_create_decorator(AccountRisk) @@ -44,64 +18,152 @@ def create_risk(data): return AccountRisk(**data) -@bulk_update_decorator(AccountRisk, update_fields=["details"]) +@bulk_update_decorator(AccountRisk, update_fields=["details", "status"]) def update_risk(risk): return risk -def check_account_secrets(accounts, assets): - now = timezone.now().isoformat() - risks = [] - tmpl = "Check account %s: %s" - summary = defaultdict(int) - result = defaultdict(list) - summary["accounts"] = len(accounts) - summary["assets"] = len(assets) +class BaseCheckHandler: + risk = '' - for account in accounts: - result_item = { - "asset": str(account.asset), - "username": account.username, - } + def __init__(self, assets): + self.assets = assets + + def check(self, account): + pass + + def clean(self): + pass + + +class CheckSecretHandler(BaseCheckHandler): + risk = RiskChoice.weak_password + + @staticmethod + def is_weak_password(password): + # 判断密码长度 + if len(password) < 8: + return True + + # 判断是否只有一种字符类型 + if password.isdigit() or password.isalpha(): + return True + + # 判断是否只包含数字或字母 + if password.islower() or password.isupper(): + return True + + # 判断是否包含常见弱密码 + common_passwords = ["123456", "password", "12345678", "qwerty", "abc123"] + if password.lower() in common_passwords: + return True + + # 正则表达式判断字符多样性（数字、字母、特殊字符） + if ( + not re.search(r"[A-Za-z]", password) + or not re.search(r"[0-9]", password) + or not re.search(r"[\W_]", password) + ): + return True + return False + + def check(self, account): if not account.secret: - print(tmpl % (account, "no secret")) - summary["no_secret"] += 1 - result["no_secret"].append(result_item) - continue + return False + return self.is_weak_password(account.secret) - if is_weak_password(account.secret): - print(tmpl % (account, color_fmt("weak", "red"))) - summary[RiskChoice.weak_password] += 1 - result[RiskChoice.weak_password].append(result_item) - risks.append( - { - "account": account, - "risk": RiskChoice.weak_password, - } - ) - else: - summary["ok"] += 1 - result["ok"].append(result_item) - print(tmpl % (account, color_fmt("ok", "green"))) - origin_risks = AccountRisk.objects.filter(asset__in=assets) - origin_risks_dict = {f"{r.asset_id}_{r.username}_{r.risk}": r for r in origin_risks} +class CheckRepeatHandler(BaseCheckHandler): + risk = RiskChoice.repeated_password - for d in risks: - key = f'{d["account"].asset_id}_{d["account"].username}_{d["risk"]}' - origin_risk = origin_risks_dict.get(key) + def __init__(self, assets): + super().__init__(assets) + self.path, self.conn, self.cursor = self.init_repeat_check_db() + self.add_password_for_check_repeat() - if origin_risk: - origin_risk.details.append({"datetime": now, 'type': 'refind'}) - update_risk(origin_risk) - else: - create_risk({ - "asset": d["account"].asset, - "username": d["account"].username, - "risk": d["risk"], - "details": [{"datetime": now, 'type': 'init'}], - }) - return summary, result + @staticmethod + def init_repeat_check_db(): + path = os.path.join('/tmp', 'accounts_' + str(uuid.uuid4()) + '.db') + sql = """ + CREATE TABLE IF NOT EXISTS accounts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + digest CHAR(32) + ) + """ + index = "CREATE INDEX IF NOT EXISTS idx_digest ON accounts(digest)" + conn = sqlite3.connect(path) + cursor = conn.cursor() + cursor.execute(sql) + cursor.execute(index) + return path, conn, cursor + + def check(self, account): + if not account.secret: + return False + + digest = self.digest(account.secret) + sql = 'SELECT COUNT(*) FROM accounts WHERE digest = ?' + self.cursor.execute(sql, [digest]) + result = self.cursor.fetchone() + if not result: + return False + return result[0] > 1 + + @staticmethod + def digest(secret): + return hashlib.md5(secret.encode()).hexdigest() + + def add_password_for_check_repeat(self): + accounts = Account.objects.all().only('id', '_secret', 'secret_type') + sql = "INSERT INTO accounts (digest) VALUES (?)" + + for account in accounts: + secret = account.secret + if not secret: + continue + digest = self.digest(secret) + self.cursor.execute(sql, [digest]) + self.conn.commit() + + def clean(self): + self.cursor.close() + self.conn.close() + os.remove(self.path) + + +class CheckLeakHandler(BaseCheckHandler): + risk = RiskChoice.leaked_password + + def __init__(self, *args): + super().__init__(*args) + self.conn, self.cursor = self.init_leak_password_db() + + @staticmethod + def init_leak_password_db(): + db_path = os.path.join( + settings.APPS_DIR, 'accounts', 'automations', + 'check_account', 'leak_passwords.db' + ) + + if settings.LEAK_PASSWORD_DB_PATH and os.path.isfile(settings.LEAK_PASSWORD_DB_PATH): + db_path = settings.LEAK_PASSWORD_DB_PATH + + db_conn = sqlite3.connect(db_path) + db_cursor = db_conn.cursor() + return db_conn, db_cursor + + def check(self, account): + if not account.secret: + return False + + sql = 'SELECT 1 FROM leak_passwords WHERE password = ? LIMIT 1' + self.cursor.execute(sql, (account.secret,)) + leak = self.cursor.fetchone() is not None + return leak + + def clean(self): + self.cursor.close() + self.conn.close() class CheckAccountManager(BaseManager): @@ -110,54 +172,40 @@ class CheckAccountManager(BaseManager): def __init__(self, execution): super().__init__(execution) - self.accounts = [] self.assets = [] - self.global_origin_risks_dict = dict() - self.db_conn = None - self.db_cursor = None + self.batch_risks = [] + self.handlers = [] - def init_leak_password_db(self): - default_path = os.path.join( - settings.APPS_DIR, 'accounts', 'automations', 'check_account' - ) - create_table = ''' - CREATE TABLE IF NOT EXISTS accounts ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - asset_id TEXT, asset_name TEXT, name TEXT, - username TEXT, password TEXT - ) - ''' - if (settings.LEAK_PASSWORD_DB_PATH - and os.path.exists(settings.LEAK_PASSWORD_DB_PATH)): - db_path = settings.LEAK_PASSWORD_DB_PATH - else: - db_path = os.path.join(default_path, 'leak_passwords.db') + def add_risk(self, risk, account): + self.summary[risk] += 1 + self.result[risk].append({ + 'asset': str(account.asset), 'username': account.username, + }) + risk_obj = {'account': account, 'risk': risk} + self.batch_risks.append(risk_obj) - self.db_conn = sqlite3.connect(db_path) - self.db_cursor = self.db_conn.cursor() - self.db_cursor.execute(create_table) + def commit_risks(self, assets): + account_risks = AccountRisk.objects.filter(asset__in=assets) + ori_risk_map = {} - def drop_account_table(self): - sql = 'DROP TABLE IF EXISTS accounts' - self.db_cursor.execute(sql) - self.db_conn.commit() + for risk in account_risks: + key = f'{risk.asset_id}_{risk.username}_{risk.risk}' + ori_risk_map[key] = risk - def close_db(self): - try: - self.db_cursor.close() - self.db_conn.close() - except Exception: # noqa - pass - - @staticmethod - def create_or_update_risk(risks, origin_risks_dict): now = timezone.now().isoformat() - for d in risks: + for d in self.batch_risks: key = f'{d["account"].asset_id}_{d["account"].username}_{d["risk"]}' - origin_risk = origin_risks_dict.get(key) + origin_risk = ori_risk_map.get(key) - if origin_risk: - origin_risk.details.append({"datetime": now, 'type': 'refind'}) + if origin_risk and origin_risk.status != ConfirmOrIgnore.pending: + details = origin_risk.details or [] + details.append({"datetime": now, 'type': 'refind'}) + + if len(details) > 10: + details = [*details[:5], *details[-5:]] + + origin_risk.details = details + origin_risk.status = ConfirmOrIgnore.pending update_risk(origin_risk) else: create_risk({ @@ -167,75 +215,8 @@ class CheckAccountManager(BaseManager): "details": [{"datetime": now, 'type': 'init'}], }) - def is_leak_password(self, password): - sql = 'SELECT 1 FROM leak_passwords WHERE password = ? LIMIT 1' - self.db_cursor.execute(sql, (password,)) - return self.db_cursor.fetchone() is not None - - def check_account_secrets(self, accounts, assets): - risks = [] - for account in accounts: - if not account.secret: - print(self.tmpl % (account, "no secret")) - self.risk_record('no_secret', account) - continue - - if is_weak_password(account.secret): - key = RiskChoice.weak_password - print(self.tmpl % (account, color_fmt(key.value, "red"))) - risks.append(self.risk_record(key, account)) - elif self.is_leak_password(account.secret): - key = RiskChoice.leaked_password - print(self.tmpl % (account, color_fmt(key.value, "red"))) - risks.append(self.risk_record(key, account)) - else: - sql = ("INSERT INTO accounts (name, username, password, asset_id, asset_name) " - "VALUES (?, ?, ?, ?, ?)") - self.db_cursor.execute( - sql, [ - account.name, account.username, account.secret, - str(account.asset_id), account.asset.name - ] - ) - self.db_conn.commit() - - origin_risks = AccountRisk.objects.filter(asset__in=assets) - origin_risks_dict = {f"{r.asset_id}_{r.username}_{r.risk}": r for r in origin_risks} - self.global_origin_risks_dict.update(origin_risks_dict) - self.create_or_update_risk(risks, origin_risks_dict) - - def risk_record(self, key, account): - self.summary[key] += 1 - self.result[key].append({ - 'asset': str(account.asset), 'username': account.username, - }) - return {'account': account, 'risk': key} - - def check_repeat_secrets(self): - risks = [] - sql = ''' - SELECT name, username, password, asset_id, asset_name, - CASE WHEN password IN ( - SELECT password FROM accounts GROUP BY password HAVING COUNT(*) > 1 - ) THEN 1 ELSE 0 END AS is_repeated FROM accounts - ''' - self.db_cursor.execute(sql) - for results in self.db_cursor.fetchall(): - name, username, *_, asset_id, asset_name, is_repeat = results - account = Account(asset_id=asset_id, username=username, name=name) - account_display = f'{name}({asset_name})' - if is_repeat: - key = RiskChoice.repeated_password - print(self.tmpl % (account_display, color_fmt(key.value, "red"))) - risks.append(self.risk_record(key, account)) - else: - key = 'ok' - print(self.tmpl % (account_display, color_fmt("ok", "green"))) - self.risk_record(key, account) - self.create_or_update_risk(risks, self.global_origin_risks_dict) - def pre_run(self): - self.init_leak_password_db() + super().pre_run() self.assets = self.execution.get_all_assets() self.execution.date_start = timezone.now() self.execution.save(update_fields=["date_start"]) @@ -243,22 +224,37 @@ class CheckAccountManager(BaseManager): def do_run(self, *args, **kwargs): for engine in self.execution.snapshot.get("engines", []): if engine == "check_account_secret": - batch_handle = self.check_account_secrets - global_handle = self.check_repeat_secrets + handler = CheckSecretHandler(self.assets) + elif engine == "check_account_repeat": + handler = CheckRepeatHandler(self.assets) + elif engine == "check_account_leak": + handler = CheckLeakHandler(self.assets) else: + print("Unknown engine: {}".format(engine)) continue + self.handlers.append(handler) + + print("Engine: {}".format(handler.__class__.__name__)) for i in range(0, len(self.assets), self.batch_size): _assets = self.assets[i: i + self.batch_size] accounts = Account.objects.filter(asset__in=_assets) - batch_handle(accounts, _assets) - global_handle() + print("Start to check accounts: {}".format(len(accounts))) + + for account in accounts: + error = handler.check(account) + print("Check: {} => {}".format(account, error)) + if not error: + continue + self.add_risk(handler.risk, account) + + self.commit_risks(_assets) def post_run(self): super().post_run() - self.drop_account_table() - self.close_db() + for handler in self.handlers: + handler.clean() def get_report_subject(self): return "Check account report of %s" % self.execution.id @@ -268,16 +264,16 @@ class CheckAccountManager(BaseManager): def print_summary(self): tmpl = ( - "\n---\nSummary: \nok: %s, weak password: %s, leaked password: %s, " - "repeated password: %s, no secret: %s, using time: %ss" - % ( - self.summary["ok"], - self.summary[RiskChoice.weak_password], - self.summary[RiskChoice.leaked_password], - self.summary[RiskChoice.repeated_password], + "\n---\nSummary: \nok: %s, weak password: %s, leaked password: %s, " + "repeated password: %s, no secret: %s, using time: %ss" + % ( + self.summary["ok"], + self.summary[RiskChoice.weak_password], + self.summary[RiskChoice.leaked_password], + self.summary[RiskChoice.repeated_password], - self.summary["no_secret"], - int(self.duration), - ) + self.summary["no_secret"], + int(self.duration), + ) ) print(tmpl) diff --git a/apps/assets/automations/base/manager.py b/apps/assets/automations/base/manager.py index d182eb88a..a8d630f57 100644 --- a/apps/assets/automations/base/manager.py +++ b/apps/assets/automations/base/manager.py @@ -1,5 +1,6 @@ import hashlib import json +import logging import os import shutil import time @@ -143,7 +144,7 @@ class BaseManager: recipients = self.execution.recipients if not recipients: return - print("Send report to: ", ",".join([str(u) for u in recipients])) + print("Send report to: ", ",".join([str(u) for u in recipients])) report = self.gen_report() report = transform(report) @@ -166,7 +167,8 @@ class BaseManager: self.pre_run() try: self.do_run(*args, **kwargs) - except: + except Exception as e: + logging.exception(e) self.status = 'error' finally: self.post_run() @@ -337,10 +339,10 @@ class PlaybookPrepareMixin: method_attr = "{}_method".format(method_type) method_enabled = ( - automation - and getattr(automation, enabled_attr) - and getattr(automation, method_attr) - and getattr(automation, method_attr) in self.method_id_meta_mapper + automation + and getattr(automation, enabled_attr) + and getattr(automation, method_attr) + and getattr(automation, method_attr) in self.method_id_meta_mapper ) if not method_enabled: @@ -427,7 +429,7 @@ class BasePlaybookManager(PlaybookPrepareMixin, BaseManager): # 避免一个任务太大，分批执行 assets_bulked = [ - assets[i : i + self.bulk_size] + assets[i: i + self.bulk_size] for i in range(0, len(assets), self.bulk_size) ] for i, _assets in enumerate(assets_bulked, start=1):