perf: 修改一些 adhoc 任务

This commit is contained in:
ibuler 2024-12-10 15:44:16 +08:00
parent 3d5b6376e8
commit c95c3099b7
5 changed files with 82 additions and 50 deletions

View File

@ -1,9 +1,12 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
# #
from django.db.models import Q, Count from django.db.models import Q, Count
from django.http import HttpResponse
from rest_framework.decorators import action from rest_framework.decorators import action
from rest_framework.exceptions import MethodNotAllowed from rest_framework.exceptions import MethodNotAllowed
from operator import itemgetter from operator import itemgetter
from django.shortcuts import get_object_or_404
from django.utils import timezone
from rest_framework.response import Response from rest_framework.response import Response
@ -14,7 +17,9 @@ from accounts.models import (
AccountRisk, AccountRisk,
RiskChoice, RiskChoice,
CheckAccountEngine, CheckAccountEngine,
AutomationExecution,
) )
from assets.models import Asset
from common.api import JMSModelViewSet from common.api import JMSModelViewSet
from common.utils import many_get from common.utils import many_get
from orgs.mixins.api import OrgBulkModelViewSet from orgs.mixins.api import OrgBulkModelViewSet
@ -42,6 +47,7 @@ class CheckAccountExecutionViewSet(AutomationExecutionViewSet):
("list", "accounts.view_checkaccountexecution"), ("list", "accounts.view_checkaccountexecution"),
("retrieve", "accounts.view_checkaccountsexecution"), ("retrieve", "accounts.view_checkaccountsexecution"),
("create", "accounts.add_checkaccountexecution"), ("create", "accounts.add_checkaccountexecution"),
("adhoc", "accounts.add_checkaccountexecution"),
("report", "accounts.view_checkaccountsexecution"), ("report", "accounts.view_checkaccountsexecution"),
) )
ordering = ("-date_created",) ordering = ("-date_created",)
@ -52,6 +58,26 @@ class CheckAccountExecutionViewSet(AutomationExecutionViewSet):
queryset = queryset.filter(automation__type=self.tp) queryset = queryset.filter(automation__type=self.tp)
return queryset return queryset
@action(methods=["get"], detail=False, url_path="adhoc")
def adhoc(self, request, *args, **kwargs):
asset_id = request.query_params.get("asset_id")
if not asset_id:
return Response(status=400, data={"asset_id": "This field is required."})
get_object_or_404(Asset, pk=asset_id)
execution = AutomationExecution()
execution.snapshot = {
"assets": [asset_id],
"nodes": [],
"type": AutomationTypes.check_account,
"engines": ["check_account_secret"],
"name": "Check asset risk: {} {}".format(asset_id, timezone.now()),
}
execution.save()
execution.start()
report = execution.manager.gen_report()
return HttpResponse(report)
class AccountRiskViewSet(OrgBulkModelViewSet): class AccountRiskViewSet(OrgBulkModelViewSet):
model = AccountRisk model = AccountRisk
@ -99,7 +125,9 @@ class AccountRiskViewSet(OrgBulkModelViewSet):
s = self.get_serializer(data=request.data) s = self.get_serializer(data=request.data)
s.is_valid(raise_exception=True) s.is_valid(raise_exception=True)
asset, username, act, risk = many_get(s.validated_data, ("asset", "username", "action", "risk")) asset, username, act, risk = many_get(
s.validated_data, ("asset", "username", "action", "risk")
)
handler = RiskHandler(asset=asset, username=username, request=self.request) handler = RiskHandler(asset=asset, username=username, request=self.request)
data = handler.handle(act, risk) data = handler.handle(act, risk)
if not data: if not data:

View File

@ -36,6 +36,7 @@ class GatherAccountsExecutionViewSet(AutomationExecutionViewSet):
("list", "accounts.view_gatheraccountsexecution"), ("list", "accounts.view_gatheraccountsexecution"),
("retrieve", "accounts.view_gatheraccountsexecution"), ("retrieve", "accounts.view_gatheraccountsexecution"),
("create", "accounts.add_gatheraccountsexecution"), ("create", "accounts.add_gatheraccountsexecution"),
("adhoc", "accounts.add_gatheraccountsexecution"),
("report", "accounts.view_gatheraccountsexecution"), ("report", "accounts.view_gatheraccountsexecution"),
) )
@ -46,6 +47,27 @@ class GatherAccountsExecutionViewSet(AutomationExecutionViewSet):
queryset = queryset.filter(automation__type=self.tp) queryset = queryset.filter(automation__type=self.tp)
return queryset return queryset
@action(methods=["get"], detail=False, url_path="adhoc")
def adhoc(self, request, *args, **kwargs):
asset_id = request.query_params.get("asset_id")
if not asset_id:
return Response(status=400, data={"asset_id": "This field is required."})
get_object_or_404(Asset, pk=asset_id)
execution = AutomationExecution()
execution.snapshot = {
"assets": [asset_id],
"nodes": [],
"type": "gather_accounts",
"is_sync_account": False,
"check_risk": True,
"name": "Adhoc gather accounts: {}".format(asset_id),
}
execution.save()
execution.start()
report = execution.manager.gen_report()
return HttpResponse(report)
class GatheredAccountViewSet(OrgBulkModelViewSet): class GatheredAccountViewSet(OrgBulkModelViewSet):
model = GatheredAccount model = GatheredAccount
@ -58,7 +80,6 @@ class GatheredAccountViewSet(OrgBulkModelViewSet):
} }
rbac_perms = { rbac_perms = {
"sync_accounts": "assets.add_gatheredaccount", "sync_accounts": "assets.add_gatheredaccount",
"discover": "assets.add_gatheredaccount",
"status": "assets.change_gatheredaccount", "status": "assets.change_gatheredaccount",
} }
@ -81,24 +102,3 @@ class GatheredAccountViewSet(OrgBulkModelViewSet):
handler = RiskHandler(asset, username, request=self.request) handler = RiskHandler(asset, username, request=self.request)
handler.handle_delete_remote() handler.handle_delete_remote()
return Response(status=status.HTTP_200_OK) return Response(status=status.HTTP_200_OK)
@action(methods=["get"], detail=False, url_path="discover")
def discover(self, request, *args, **kwargs):
asset_id = request.query_params.get("asset_id")
if not asset_id:
return Response(status=400, data={"asset_id": "This field is required."})
get_object_or_404(Asset, pk=asset_id)
execution = AutomationExecution()
execution.snapshot = {
"assets": [asset_id],
"nodes": [],
"type": "gather_accounts",
"is_sync_account": False,
"check_risk": True,
"name": "Adhoc gather accounts: {}".format(asset_id),
}
execution.save()
execution.start()
report = execution.manager.gen_report()
return HttpResponse(report)

View File

@ -123,7 +123,7 @@ class CheckAccountManager(BaseManager):
continue continue
for i in range(0, len(self.assets), self.batch_size): for i in range(0, len(self.assets), self.batch_size):
_assets = self.assets[i : i + self.batch_size] _assets = self.assets[i: i + self.batch_size]
accounts = Account.objects.filter(asset__in=_assets) accounts = Account.objects.filter(asset__in=_assets)
summary, result = handle(accounts, _assets) summary, result = handle(accounts, _assets)

View File

@ -84,6 +84,8 @@
{% endfor %} {% endfor %}
</tbody> </tbody>
</table> </table>
{% else %}
<p>{% trans 'No weak password' %}</p>
{% endif %} {% endif %}
</div> </div>

View File

@ -68,6 +68,8 @@
{% endfor %} {% endfor %}
</tbody> </tbody>
</table> </table>
{% else %}
<p>{% trans 'No new accounts found' %}</p>
{% endif %} {% endif %}
</div> </div>
<div class='result'> <div class='result'>
@ -96,32 +98,32 @@
</div> </div>
<div class='result'> {#<div class='result'>#}
<p>{% trans 'New found risks' %}: {{ summary.new_risks }}</p> {# <p>{% trans 'New found risks' %}: {{ summary.new_risks }}</p>#}
{% if summary.new_risks %} {# {% if summary.new_risks %}#}
<table> {# <table>#}
<caption></caption> {# <caption></caption>#}
<thead> {# <thead>#}
<tr> {# <tr>#}
<th>{% trans 'No.' %}</th> {# <th>{% trans 'No.' %}</th>#}
<th>{% trans 'Asset' %}</th> {# <th>{% trans 'Asset' %}</th>#}
<th>{% trans 'Username' %}</th> {# <th>{% trans 'Username' %}</th>#}
<th>{% trans 'Result' %}</th> {# <th>{% trans 'Result' %}</th>#}
</tr> {# </tr>#}
</thead> {# </thead>#}
<tbody> {# <tbody>#}
{% for risk in result.risks %} {# {% for risk in result.risks %}#}
<tr> {# <tr>#}
<td>{{ forloop.counter }}</td> {# <td>{{ forloop.counter }}</td>#}
<td>{{ risk.asset }}</td> {# <td>{{ risk.asset }}</td>#}
<td>{{ risk.username }}</td> {# <td>{{ risk.username }}</td>#}
<td>{{ risk.risk }}</td> {# <td>{{ risk.risk }}</td>#}
</tr> {# </tr>#}
{% endfor %} {# {% endfor %}#}
</tbody> {# </tbody>#}
</table> {# </table>#}
{% endif %} {# {% endif %}#}
</div> {#</div>#}
<style> <style>
table { table {