zkyuan/src/script/start_collect_mtask_result_webshell.py

#-*-coding:utf-8 -*-
import json
import os,sys
import django

sys.path.append(os.path.join("..",os.path.dirname(__file__)))
os.environ['DJANGO_SETTINGS_MODULE'] = 'settings'
django.setup()

import common.models as cm
import webray_service.control_wshellpoc as cwsp


def get_available_tasks():
    """
    """
    wrtask = cm.WebRayTasks.objects.filter(type=3)
    return wrtask


def collect_results_webshell():
    tasks = get_available_tasks()
    for task in tasks:
        siteid = task.siteid
        task_id = str(task.id)
        website_id = task.website_id
        wstrobj = cm.WebShellTasksResult.objects.filter(task_id=task.id).order_by("-jobid").first()
        jobid = wstrobj.jobid if wstrobj else -1
        if task.excute_type == "once":
            if jobid > -1:
                continue
            else:
                jobid = 0
        else:
            jobid = jobid + 1
        wajob_res = cwsp.get_webshellpoc_results([{"task_id":task_id,"job_id":jobid}])
        print wajob_res,2222222222
        if wajob_res.get("message") == "ok":
            if wajob_res.get("data"):
                data = wajob_res.get("data")[0]
                trobj,flag = cm.WebShellTasksResult.objects.get_or_create(
                        task_type = 4,
                        task_id = task.id, 
                        website_id = website_id, 
                        taskid = task.id,
                        jobid = jobid,
                        task_status = data.get("task_status"),
                        starttime = data.get("task_begintime"),
                        endtime = data.get("task_endtime"),
                        task_executing_engine_id = data.get("task_executing_engine_id"),
                        task_executing_enginevendor_id = data.get("task_executing_enginevendor_id"),
                        task_result = json.dumps(data.get("task_result")),
                        total = len(data.get("task_result",[]))
                    )
                trobj.save()
                #生成webshell告警事件
                if json.loads(trobj.task_result):
                    cm.WarningEvent.objects.create(**{"event_type":"webshell",
                            "task_type":1,
                            "task_id":task_id,
                            "taskid":task_id,
                            "jobid":jobid,
                            "taskname":task.name,
                            "domain":task.site.split("//")[1].split(":")[0],
                            "label":task.website_name,
                            "organization":cm.WebSite.objects.filter(id=website_id).first().organization_name,
                            "result":trobj.task_result
                            })
        
def main():
    collect_results_webshell()

if __name__ == "__main__":
    main()