更改数据库

8 files changed, 126 insertions, 41 deletions

diff --git a/agent/apps/target_gz.py b/agent/apps/target_gz.py
index a924401..d3287d8 100644
--- a/agent/apps/target_gz.py
+++ b/agent/apps/target_gz.py
@@ -20,8 +20,5 @@ def execute_command(IP):
                                 encoding='utf-8')
 
         output = proc.communicate()[0].strip()
-        # TODO 报错的情况
-        if "udp" in output:
-            output = None
-
+    
     return output
\ No newline at end of file
diff --git a/server/app.py b/server/app.py
index f70b8cb..43dca54 100644
--- a/server/app.py
+++ b/server/app.py
@@ -31,6 +31,9 @@ app.config[
     "SQLALCHEMY_DATABASE_URI"] = f"mysql+pymysql://root:{settings.MYSQL_PAWD}@{settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATADB}"
 # app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///project.db"
 app.config['SECRET_KEY'] = 'your-secret-key'
+
+app.debug = False
+
 db.init_app(app)
 migrate = Migrate(app, db)
 
diff --git a/server/apps/policy.py b/server/apps/policy.py
index 8a54817..0729130 100644
--- a/server/apps/policy.py
+++ b/server/apps/policy.py
@@ -9,32 +9,17 @@ from model import Policy, Target, TaskPolicy
 from .target import target_GZ
 
 def chiose_policy(task):
-    if task.target_domain is not None:
-        res = None
-        # TODO:DoH处理
-        pass
-    else:
-        ip_version = is_ipaddress(task.target_ip)
-        if ip_version == 4 :
-            res = db.session.query(Target).filter(Target.addrv4==task.target_ip).first()
-        elif ip_version == 6:
-            res = db.session.query(Target).filter(Target.addrv6==task.target_ip).first()
-
+    target = task.target
     # 目标支持协议范围
     proto = []
 
-    # 不存在该目标的相关记录
-    if not res:
-        target_GZ(task.target_ip)
-        proto = ["IPv6"]
-
-    if res.ipv6:
+    if target.ipv6:
         proto.append("IPv6")
-    if res.doh:
+    if target.doh:
         proto.append("DOH")
-    if res.dot:
+    if target.dot:
         proto.append("DOT")
-    if res.dnssec:
+    if target.dnssec:
         proto.append("DNSSEC")
 
     # 可选策略范围 policy
@@ -82,4 +67,14 @@ def copy_task_policy(task):
     db.session.add(task_policy)
     db.session.commit()
 
-    return task_policy
\ No newline at end of file
+    return task_policy
+
+def parameters2polocy(task_policy):
+    commmand = "aaaaaa"
+    return commmand
+
+def generate_parameters(task_policy):
+    pass
+
+def adjust_parameters(task_policy):
+    pass
diff --git a/server/apps/target.py b/server/apps/target.py
index 6b48e15..97722ac 100644
--- a/server/apps/target.py
+++ b/server/apps/target.py
@@ -14,7 +14,7 @@ from apiflask.validators import OneOf
 from requests.exceptions import Timeout
 from sqlalchemy import distinct, func, case
 
-from apps.util import debug
+from apps.util import debug, is_ipaddress
 from exts import db
 from model import Target, Task, Agent
 
@@ -510,8 +510,14 @@ def map_info(query_data):
     "code": Integer(),
     "dataObject": List(Nested(TargetSchema())),
 })
-def target_GZ(query_data):
-    IP_addr = query_data["ip"]
+def target_GZ_API(query_data):
+    target_GZ(query_data["ip"])
+    return {
+        'code': 200,
+        'MSG': "success"
+    }
+
+def target_GZ(IP_addr):
     # 获取
     csgz = db.session.query(Agent).filter_by(agent_type = 'gjst').all()
     csgz = random.choice(csgz)
@@ -524,10 +530,10 @@ def target_GZ(query_data):
     data = response.json()
 
     # 存数据库
-    tagret = Target(
+    target = Target(
         addrv4 = IP_addr,
         addrv6 = False,
-        ipv6 = False,
+        ipv6 = (6 == is_ipaddress(IP_addr)),
         dnssec = json.loads(protect.text)['dnssec_enabled'],
         dot = False,
         doh = False,
@@ -535,15 +541,18 @@ def target_GZ(query_data):
         isp = data.get('org'),
         lat = data.get('loc').split(',')[0],
         lng = data.get('loc').split(',')[1],
-        protect = json.loads(protect.text),
+        protect = protect.text,
         doh_domain = None
     )
-    existing_obj = db.session.query(Target).filter_by(addrv4 = tagret.addrv4).first()
+    if 6 == is_ipaddress(IP_addr):
+        target.addrv6 = IP_addr
+        target.ipv6 = True
+    elif 4== is_ipaddress(IP_addr):
+        target.addrv4 = IP_addr
+
+    existing_obj = db.session.query(Target).filter_by(addrv4 = target.addrv4).first()
     if existing_obj:
         db.session.delete(existing_obj)
-    db.session.add(tagret)
+    db.session.add(target)
     db.session.commit()
-    return {
-        'code': 200,
-        'dataObject': tagret
-    }
+    return target
\ No newline at end of file
diff --git a/server/apps/task.py b/server/apps/task.py
index 4478586..bd96646 100644
--- a/server/apps/task.py
+++ b/server/apps/task.py
@@ -109,7 +109,7 @@ def make_task(json_data):
         task_name=json_data.get("name"),
         agent_id=json_data.get("agent"),
         # created_by = "Admin",
-        target_ip=str(json_data.get("target")),
+        # target_ip=str(json_data.get("target")),
         ptype = json_data["policy"],
         status="working" if json_data.get("run_flag") == "now" else "stop",
         policy_delay=json_data.get("policy_time"),
@@ -126,6 +126,19 @@ def make_task(json_data):
     if task.target_rr is not None:
         task.target_rtype, task.target_rr = task.target_rr.split(" ")
 
+    # 关联目标
+    Ip = str(json_data.get("target"))
+    ip_version = is_ipaddress(Ip)
+    target = None
+    if ip_version == 4 :
+        target = db.session.query(Target).filter(Target.addrv4==Ip).first()
+    elif ip_version == 6:
+        target = db.session.query(Target).filter(Target.addrv6==Ip).first()
+    # 不存在该目标的相关记录，自动探测并存入数据库
+    if not target:
+        target = target_GZ(Ip)
+    task.target = target
+
     # 查找所有在线的用于状态感知的agent
     agents = db.session.query(Agent).filter(and_(Agent.status == True, Agent.agent_type == 'ztgz')).all()
     agents = random.sample(agents, min(10, len(agents)))
@@ -227,7 +240,7 @@ def tasks_state(query_data):
     for task in tasks:
         task_r = {}
         task_r["id"] = task.task_id
-        task_r["target"] = task.target_ip
+        task_r["target"] = f"{task.target.addrv4}|{task.target.addrv6}"
         task_r["name"] = task.task_name
         task_r["agent"] = task.agent_id
         task_r["target_domain"] = task.target_domain,
@@ -305,7 +318,7 @@ def taskpolicy_log(query_data):
 
     log_data = [{
         "time": log.created_time,
-        "targetip": log.task_policy.task.target_ip,
+        "targetip": f"{log.task_policy.task.target.addrv4}|{log.task_policy.task.target.addrv4}",
         "ip": log.task_policy.task.agent.ipaddr,
         "level": log.tlog_level,
         "info": log.tlog_info
diff --git a/server/migrations/versions/176af7a75d22_.py b/server/migrations/versions/176af7a75d22_.py
new file mode 100644
index 0000000..2852f27
--- /dev/null
+++ b/server/migrations/versions/176af7a75d22_.py
@@ -0,0 +1,34 @@
+"""empty message
+
+Revision ID: 176af7a75d22
+Revises: 4b6338477842
+Create Date: 2024-09-01 19:14:21.739644
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '176af7a75d22'
+down_revision = '4b6338477842'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('TASK', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('target_id', sa.Integer(), nullable=True))
+        batch_op.create_foreign_key(None, 'TARGETDATA', ['target_id'], ['target_id'])
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('TASK', schema=None) as batch_op:
+        batch_op.drop_constraint(None, type_='foreignkey')
+        batch_op.drop_column('target_id')
+
+    # ### end Alembic commands ###
diff --git a/server/migrations/versions/4b6338477842_.py b/server/migrations/versions/4b6338477842_.py
new file mode 100644
index 0000000..ec7ab23
--- /dev/null
+++ b/server/migrations/versions/4b6338477842_.py
@@ -0,0 +1,32 @@
+"""empty message
+
+Revision ID: 4b6338477842
+Revises: 854f2615247e
+Create Date: 2024-09-01 19:11:39.853533
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision = '4b6338477842'
+down_revision = '854f2615247e'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('TASK', schema=None) as batch_op:
+        batch_op.drop_column('target_ip')
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('TASK', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('target_ip', mysql.VARCHAR(collation='utf8_unicode_ci', length=255), nullable=True))
+
+    # ### end Alembic commands ###
diff --git a/server/model.py b/server/model.py
index 1c56092..7f0c214 100644
--- a/server/model.py
+++ b/server/model.py
@@ -33,7 +33,8 @@ class Task(db.Model):
     task_name = db.Column(db.String(255))
     agent_id = db.Column(db.String(255), db.ForeignKey(f'{MYSQL_TAB_AGENT}.agent_id'))
     created_by = db.Column(db.String(255), db.ForeignKey(f'{MYSQL_TAB_USER}.user_id'))
-    target_ip = db.Column(db.String(255))
+    # target_ip = db.Column(db.String(255))
+    target_id = db.Column(db.Integer, db.ForeignKey(f'{MYSQL_TAB_TARGETDATA}.target_id'))
     created_time = db.Column(db.DateTime, server_default=db.func.now())
     ptype = db.Column(db.String(255))
     status = db.Column(db.String(255))
@@ -144,6 +145,7 @@ class Target(db.Model):
     })
     doh_domain = db.Column(db.String(255))
 
+    tasks = db.relationship('Task', backref='target')
 
 # 对单个模型操作
 def model_to_dict(model, fields=None, exchange_fields=None):