perf: 重构扫描

pull/14788/head
ibuler 2025-01-08 19:13:13 +08:00
parent 08bf8decdd
commit 49f34b8124
4 changed files with 229 additions and 216 deletions

View File

@ -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()

View File

@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:8526225e13163cc92535a0ed3ef1c12cfde471173394504aea2b0583fe68efb5
oid sha256:37be0604185cf3189de457dc7339a461052d96582096beb4fdf7b0fcb2c290a8
size 98037760

View File

@ -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)

View File

@ -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):