admin.py

import threading
from django.db import transaction
from django.contrib import admin, messages
from VulnerabilityMonitor.models import VulnerabilityMonitorTask, VulnerbilityMonitorResult
from VulnerabilityMonitor.views import start_monitor
from Configuration.models import Configuration

# Register your models here.
admin.site.site_header = '阿波罗自动化攻击评估系统'  # 设置header
admin.site.site_title = '阿波罗自动化攻击评估系统'  # 设置title

@admin.register(VulnerabilityMonitorTask)
class VulnerabilityMonitorTaskAdmin(admin.ModelAdmin):
    list_display = ['name', 'keyword', 'timestamp', 'change']
    list_filter = []
    search_fields = ['name']
    ordering = ["-id"]
    date_hierarchy = 'timestamp'

    @transaction.atomic
    def scan(self, request, queryset):
        try:
            Configuration.objects.filter(name="3").values_list("value")[0][0]
        except Exception as execption:
            print(execption)
            messages.add_message(request, messages.SUCCESS, '请配置Github接口Token')
            return
        work_ids = None
        for item in request.POST.lists():
            if item[0] == "_selected_action":
                work_ids = item[1]
        if isinstance(work_ids, list):
            for work_id in work_ids:
                thread = threading.Thread(target=start_monitor, args=(work_id,))
                thread.start()
                messages.add_message(request, messages.SUCCESS, '开始监控%s' % str(work_id))
        else:
            messages.add_message(request, messages.SUCCESS, '监控异常')

    scan.short_description = "启动监控"
    scan.icon = 'fa fa-rocket'
    scan.style = 'color:white;'
    scan.type = 'danger'
    scan.confirm = '您确定要启动监控吗？'
    actions = [scan, ]


@admin.register(VulnerbilityMonitorResult)
class VulnerbilityMonitorResultAdmin(admin.ModelAdmin):
    list_display = ['username', 'taskname', 'projectname' ,'linkurl', 'state', 'timestamp', 'change']
    list_filter = ['state', 'taskname']
    search_fields = ['username', 'projectname']
    ordering = ["-id"]
    date_hierarchy = 'timestamp'

    def has_add_permission(self, request):
        return False

    def has_change_permission(self, request, obj=None):
        return True