feat: add tantivy search

.github/actions/setup-python/action.yml

@@ -0,0 +1,26 @@
+name: Setup Python
+description: Setup Python
+
+inputs:
+  python-version:
+    description: Python version
+    required: false
+    default: "3.10"
+
+runs:
+  using: "composite"
+  steps:
+    - uses: actions/setup-python@v4
+      with:
+        python-version: ${{ inputs.python-version }}
+        cache: "pip"
+
+    - name: Install dependencies
+      shell: bash
+      run: |
+        sudo apt-get install build-essential cmake ninja-build python3-dev cython3 pybind11-dev libre2-dev -y
+        python -m pip install --upgrade pip
+        pip install wheel maturin
+        pip install -r requirements.txt
+        curl -L https://github.com/HXSecurity/tantivy-py/releases/download/0.21.0/tantivy-0.20.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl -o /tmp/tantivy-0.20.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+        pip install /tmp/tantivy-0.20.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

.github/workflows/nightly.yaml

@@ -37,18 +37,11 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v2
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-          cache: "pip"
+      - name: Setup Python
+        uses: ./.github/actions/setup-python
 
       - name: Install dependencies
-        run: |
-          sudo apt-get install build-essential cmake ninja-build python3-dev cython3 pybind11-dev libre2-dev -y
-          python -m pip install --upgrade pip
-          pip install wheel
-          pip install coverage
-          pip install -r requirements.txt
+        run: pip install coverage
 
       - name: Django Unit Testing
         run: |
@@ -92,17 +85,11 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v2
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-          cache: "pip"
+      - name: Setup Python
+        uses: ./.github/actions/setup-python
 
       - name: Install dependencies
         run: |
-          sudo apt-get install build-essential cmake ninja-build python3-dev cython3 pybind11-dev libre2-dev -y
-          python -m pip install --upgrade pip
-          pip install wheel
-          pip install -r requirements.txt
           pip install schemathesis
           pip install httpx
 

.github/workflows/teststate.yml

@@ -55,17 +55,11 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v2
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-          cache: "pip"
+      - name: Setup Python
+        uses: ./.github/actions/setup-python
 
       - name: Install dependencies
         run: |
-          sudo apt-get install build-essential cmake ninja-build python3-dev cython3 pybind11-dev libre2-dev -y
-          python -m pip install --upgrade pip
-          pip install wheel
-          pip install -r requirements.txt
           pip install schemathesis
           pip install httpx
 
@@ -124,18 +118,8 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v2
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-          cache: "pip"
-
-      - name: Install dependencies
-        run: |
-          sudo apt-get install build-essential cmake ninja-build python3-dev cython3 pybind11-dev libre2-dev -y
-          python -m pip install --upgrade pip
-          pip install wheel
-          pip install -r requirements.txt
-
+      - name: Setup Python
+        uses: ./.github/actions/setup-python
       - name: Django Unit Testing
         run: |
           mypy --show-error-codes --ignore-missing-imports  --no-incremental --show-error-codes --check-untyped-defs --disable-error-code var-annotated  --disable-error-code  list-item  --disable-error-code attr-defined --disable-error-code assignment --disable-error-code misc --disable-error-code union-attr --disable-error-code index --disable-error-code call-overload  --disable-error-code dict-item  --disable-error-code truthy-function --disable-error-code operator --disable-error-code name-defined .
@@ -168,18 +152,11 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v2
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-          cache: "pip"
+      - name: Setup Python
+        uses: ./.github/actions/setup-python
 
       - name: Install dependencies
-        run: |
-          sudo apt-get install build-essential cmake ninja-build python3-dev cython3 pybind11-dev libre2-dev -y
-          python -m pip install --upgrade pip
-          pip install wheel
-          pip install -r requirements.txt
-          pip install bandit
+        run: pip install bandit
       - name: Django Unit Testing
         run: |
           bandit -iii -lll -r .
@@ -215,18 +192,11 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v2
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-          cache: "pip"
+      - name: Setup Python
+        uses: ./.github/actions/setup-python
 
       - name: Install dependencies
-        run: |
-          sudo apt-get install build-essential cmake ninja-build python3-dev cython3 pybind11-dev libre2-dev -y
-          python -m pip install --upgrade pip
-          pip install wheel
-          pip install coverage
-          pip install -r requirements.txt
+        run: pip install coverage
 
       - name: Django Unit Testing
         run: |
@@ -246,18 +216,8 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v2
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-          cache: "pip"
-
-      - name: Install dependencies
-        run: |
-          sudo apt-get install build-essential cmake ninja-build python3-dev cython3 pybind11-dev libre2-dev -y
-          python -m pip install --upgrade pip
-          pip install wheel
-          pip install coverage
-          pip install -r requirements.txt
+      - name: Setup Python
+        uses: ./.github/actions/setup-python
 
       - name: Check Schema
         run: python3 manage.py check_schema

Dockerfile

@@ -9,15 +9,16 @@ ENV TZ=Asia/Shanghai
 RUN apt-get update -y \
     && apt install -y gettext gcc make cmake libmariadb-dev curl libc6-dev libxrender1 libxtst6 libxi6 unzip cron \
     fonts-wqy-microhei vim build-essential ninja-build cython3 pybind11-dev libre2-dev locales \
-#   htop sysstat net-tools iproute2 procps lsof \
+    # htop sysstat net-tools iproute2 procps lsof \
     zip libjpeg62 \
     && sed -i '/en_US.UTF-8/s/^# //g' /etc/locale.gen && locale-gen \
     && ALIMARCH=`arch` && curl -L https://charts.dongtai.io/apk/${ALIMARCH}/wkhtmltopdf -o /usr/bin/wkhtmltopdf \
-    && chmod +x /usr/bin/wkhtmltopdf
+    && chmod +x /usr/bin/wkhtmltopdf \
+    && curl -L https://github.com/HXSecurity/tantivy-py/releases/download/0.21.0/tantivy-0.20.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl -o /tmp/tantivy-0.20.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
 
 COPY Pipfile .
 COPY Pipfile.lock .
-RUN pip install -U pip && pip install pipenv wheel && python3 -m pipenv sync --system -v
+RUN pip install -U pip && pip install pipenv wheel && python3 -m pipenv sync --system -v && pip install /tmp/tantivy-0.20.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
 
 COPY . /opt/dongtai
 WORKDIR /opt/dongtai

dongtai_common/migrations/0032_update_vul_index_task.py

@@ -0,0 +1,23 @@
+from django.db import migrations
+
+
+def update_vul_index_task(apps, schema_editor):
+    from django_celery_beat.models import IntervalSchedule, PeriodicTask
+
+    obj = IntervalSchedule.objects.create(id=8, every=30, period=IntervalSchedule.MINUTES)
+    PeriodicTask.objects.create(
+        name="update_vul_tantivy_index",
+        task="dongtai_web.aggr_vul.tasks.update_vul_tantivy_index",
+        enabled=True,
+        interval=obj,
+    )
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("dongtai_common", "0031_auto_20230926_1510"),
+    ]
+
+    operations = [
+        migrations.RunPython(update_vul_index_task),
+    ]

dongtai_common/models/vulnerablity.py

@@ -1,6 +1,8 @@
 import logging
+import os.path
 import uuid
 
+import tantivy
 from django.core.cache import cache
 from django.db import models
 from django_elasticsearch_dsl import Document, fields
@@ -15,7 +17,7 @@
 from dongtai_common.models.strategy import IastStrategyModel
 from dongtai_common.models.vul_level import IastVulLevel
 from dongtai_common.utils.settings import get_managed
-from dongtai_conf.settings import VULNERABILITY_INDEX
+from dongtai_conf.settings import TANTIVY_INDEX_PATH, VULNERABILITY_INDEX
 
 logger = logging.getLogger("dongtai-core")
 
@@ -190,3 +192,34 @@ class Django:
         auto_refresh = False
 
         ignore_signals = False
+
+
+VUL_TANTIVY_FIELDS = [
+    "id",
+    "title",
+    "project_id",
+    "project_version_id",
+    "uri",
+    "strategy_id",
+    "level_id",
+    "status_id",
+]
+
+
+def tantivy_schema() -> tantivy.Schema:
+    schema_builder = tantivy.SchemaBuilder()
+    schema_builder.add_integer_field("id", stored=True, indexed=True)
+    schema_builder.add_text_field("title", stored=True, tokenizer_name="jieba")
+    schema_builder.add_integer_field("project_id", stored=True, indexed=True)
+    schema_builder.add_integer_field("project_version_id", stored=True, indexed=True)
+    schema_builder.add_text_field("uri", stored=True, tokenizer_name="raw")
+    schema_builder.add_integer_field("strategy_id", stored=True, indexed=True)
+    schema_builder.add_integer_field("level_id", stored=True, indexed=True)
+    schema_builder.add_integer_field("status_id", stored=True, indexed=True)
+    return schema_builder.build()
+
+
+def tantivy_index() -> tantivy.Index:
+    path = os.path.join(TANTIVY_INDEX_PATH, "vulnerability_index")
+    os.makedirs(path, exist_ok=True)
+    return tantivy.Index(tantivy_schema(), path=path)

dongtai_conf/settings.py

@@ -995,6 +995,17 @@ def safe_execute(default, exception, function, *args):
 DONGTAI_REDIS_ES_UPDATE_BATCH_SIZE = 500
 DONGTAI_MAX_BATCH_TASK_CONCORRENCY = 5
 
+
+try:
+    TANTIVY_STATE = config.get("tantivy", "enable") == "true"
+except Exception:
+    TANTIVY_STATE = False
+try:
+    TANTIVY_INDEX_PATH = config.get("tantivy", "index_path")
+except Exception:
+    TANTIVY_INDEX_PATH = "/tmp/tantivy/index_path"
+
+
 ELASTICSEARCH_STATE = config.get("elastic_search", "enable") == "true"
 
 
@@ -1045,7 +1056,6 @@ def is_gevent_monkey_patched() -> bool:
 def set_asyncio_policy():
     state = is_gevent_monkey_patched()
     print(f"is in gevent patched : {state}")
-    pass
 
 
 #   disable until this package update

dongtai_engine/tasks.py

@@ -360,7 +360,6 @@ def heartbeat():
         )
         if resp.status_code == 200:
             logger.info("[dongtai_engine.tasks.heartbeat] send heartbeat data to OpenApi Service Successful.")
-            pass
         logger.info("[dongtai_engine.tasks.heartbeat] send heartbeat data to OpenApi Service Failure.")
     except Exception as e:
         logger.info(f"[dongtai_engine.tasks.heartbeat] send heartbeat data to OpenApi Service Error. reason is {e}")

dongtai_protocol/report/handler/auth_info_handler.py

@@ -58,7 +58,6 @@ def save(self):
             )
             if len(auth_model):
                 logger.info("权限已存在,忽略")
-                pass
             else:
                 logger.info("新增权限")
                 IastOverpowerUserAuth(