add

api/dictionary.py

@@ -3,75 +3,154 @@
 # @auth: rainy-autumn@outlook.com
 # @version:
 from bson import ObjectId
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, File, UploadFile
+from starlette.responses import StreamingResponse
+
 from api.users import verify_token
-from motor.motor_asyncio import AsyncIOMotorCursor
+from motor.motor_asyncio import AsyncIOMotorCursor, AsyncIOMotorGridFSBucket
 from core.db import get_mongo_db
 from core.redis_handler import refresh_config
 from loguru import logger
 router = APIRouter()
 
+# @router.get("/subdomain/data")
+# async def get_subdomain_data(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
+#     try:
+#         # Find document with name equal to "DomainDic"
+#         result = await db.config.find_one({"name": "DomainDic"})
+#         return {
+#             "code": 200,
+#             "data": {
+#                 "dict": result.get("value", '')
+#             }
+#         }
+#
+#     except Exception as e:
+#         logger.error(str(e))
+#         # Handle exceptions as needed
+#         return {"message": "error","code":500}
 @router.get("/subdomain/data")
 async def get_subdomain_data(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
     try:
-        # Find document with name equal to "DomainDic"
+        fs = AsyncIOMotorGridFSBucket(db)
-        result = await db.config.find_one({"name": "DomainDic"})
-        return {
-            "code": 200,
-            "data": {
-                "dict": result.get("value", '')
-            }
-        }
 
+        # 查找文件
+        file_doc = await fs.find({"filename": "DomainDic"}).to_list(1)
+
+        if not file_doc:
+            return {'code': 404, 'message': 'file is not found'}
+
+        file_id = file_doc[0]['_id']
+        grid_out = await fs.open_download_stream(file_id)
+
+        # 返回文件流
+        return StreamingResponse(grid_out, media_type="application/octet-stream",
+                                 headers={"Content-Disposition": f"attachment; filename=DomainDic"})
     except Exception as e:
         logger.error(str(e))
-        # Handle exceptions as needed
+
-        return {"message": "error","code":500}
 
 @router.post("/subdomain/save")
-async def save_subdomain_data(data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
+async def save_subdomain_data(file: UploadFile = File(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
     try:
-        # Update the document with name equal to "DomainDic"
+        content = await file.read()
-        result = await db.config.update_one({"name": "DomainDic"}, {"$set": {"value": data.get('dict','')}}, upsert=True)
+        fs = AsyncIOMotorGridFSBucket(db)
-        if result.modified_count > 0:
-            await refresh_config('all', 'subdomain')
-            return {"code": 200, "message": "Successfully updated DomainDic value"}
-        else:
-            return {"code": 404, "message": "DomainDic not found"}
 
+        old_file = await fs.find({'filename': 'DomainDic'}).to_list(1)
+        if old_file:
+            await fs.delete(old_file[0]['_id'])
+
+        await fs.upload_from_stream('DomainDic', content)
+        await refresh_config('all', 'subdomain')
+        return {"code": 200, "message": "upload successful"}
     except Exception as e:
         logger.error(str(e))
         # Handle exceptions as needed
         return {"message": "error", "code": 500}
+# @router.post("/subdomain/save")
+# async def save_subdomain_data(data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
+#     try:
+#         # Update the document with name equal to "DomainDic"
+#         result = await db.config.update_one({"name": "DomainDic"}, {"$set": {"value": data.get('dict','')}}, upsert=True)
+#         if result.modified_count > 0:
+#             await refresh_config('all', 'subdomain')
+#             return {"code": 200, "message": "Successfully updated DomainDic value"}
+#         else:
+#             return {"code": 404, "message": "DomainDic not found"}
+#
+#     except Exception as e:
+#         logger.error(str(e))
+#         # Handle exceptions as needed
+#         return {"message": "error", "code": 500}
+
+# @router.get("/dir/data")
+# async def get_dir_data(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
+#     try:
+#         # Find document with name equal to "DomainDic"
+#         result = await db.config.find_one({"name": "DirDic"})
+#         return {
+#             "code": 200,
+#             "data": {
+#                 "dict": result.get("value", '')
+#             }
+#         }
+#
+#     except Exception as e:
+#         logger.error(str(e))
+#         # Handle exceptions as needed
+#         return {"message": "error","code":500}
+
 
 @router.get("/dir/data")
 async def get_dir_data(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
     try:
-        # Find document with name equal to "DomainDic"
+        fs = AsyncIOMotorGridFSBucket(db)
-        result = await db.config.find_one({"name": "DirDic"})
-        return {
-            "code": 200,
-            "data": {
-                "dict": result.get("value", '')
-            }
-        }
 
+        # 查找文件
+        file_doc = await fs.find({"filename": "dirdict"}).to_list(1)
+
+        if not file_doc:
+            return {'code': 404, 'message': 'file is not found'}
+
+        file_id = file_doc[0]['_id']
+        grid_out = await fs.open_download_stream(file_id)
+
+        # 返回文件流
+        return StreamingResponse(grid_out, media_type="application/octet-stream",
+                                 headers={"Content-Disposition": f"attachment; filename=dirdict"})
     except Exception as e:
         logger.error(str(e))
-        # Handle exceptions as needed
+
-        return {"message": "error","code":500}
+# @router.post("/dir/save")
+# async def save_subdomain_data(data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
+#     try:
+#         # Update the document with name equal to "DomainDic"
+#         result = await db.config.update_one({"name": "DirDic"}, {"$set": {"value": data.get('dict','')}}, upsert=True)
+#         if result.modified_count > 0:
+#             await refresh_config('all', 'dir')
+#             return {"code": 200, "message": "Successfully updated DirDic value"}
+#         else:
+#             return {"code": 404, "message": "DirDic not found"}
+#
+#     except Exception as e:
+#         logger.error(str(e))
+#         # Handle exceptions as needed
+#         return {"message": "error", "code": 500}
+
 
 @router.post("/dir/save")
-async def save_subdomain_data(data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
+async def save_dir_data(file: UploadFile = File(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
     try:
-        # Update the document with name equal to "DomainDic"
+        content = await file.read()
-        result = await db.config.update_one({"name": "DirDic"}, {"$set": {"value": data.get('dict','')}}, upsert=True)
+        fs = AsyncIOMotorGridFSBucket(db)
-        if result.modified_count > 0:
-            await refresh_config('all', 'dir')
-            return {"code": 200, "message": "Successfully updated DirDic value"}
-        else:
-            return {"code": 404, "message": "DirDic not found"}
 
+        old_file = await fs.find({'filename': 'dirdict'}).to_list(1)
+        if old_file:
+            await fs.delete(old_file[0]['_id'])
+
+        await fs.upload_from_stream('dirdict', content)
+        await refresh_config('all', 'dir')
+        return {"code": 200, "message": "upload successful"}
     except Exception as e:
         logger.error(str(e))
         # Handle exceptions as needed

core/db.py

@@ -78,9 +78,9 @@ async def create_database():
                 {"name": "SubfinderApiConfig", 'value': subfinderApiConfig, 'type': 'subfinder'})
             await collection.insert_one(
                 {"name": "RadConfig", 'value': radConfig, 'type': 'rad'})
-            dirDict = get_dirDict()
+            # dirDict = get_dirDict()
-            await collection.insert_one(
+            # await collection.insert_one(
-                {"name": "DirDic", 'value': dirDict, 'type': 'dirDict'})
+            #     {"name": "DirDic", 'value': dirDict, 'type': 'dirDict'})
             await collection.insert_one(
                 {"name": "notification", 'dirScanNotification': True,
                  'portScanNotification': True, 'sensitiveNotification': True,
@@ -89,9 +89,9 @@ async def create_database():
                  'subdomainNotification': True,
                  'vulNotification': True,
                  'type': 'notification'})
-            domainDict = get_domainDict()
+            # domainDict = get_domainDict()
-            await collection.insert_one(
+            # await collection.insert_one(
-                {"name": "DomainDic", 'value': domainDict, 'type': 'domainDict'})
+            #     {"name": "DomainDic", 'value': domainDict, 'type': 'domainDict'})
             sensitive_data = get_sensitive()
             collection = client[DATABASE_NAME]["SensitiveRule"]
             if sensitiveList:

core/default.py

@@ -26,7 +26,7 @@ def get_domainDict():
     domainDict = ""
     try:
         # 尝试打开文件并读取内容
-        with open(os.path.join(combined_directory, "domainDict"), "r") as file:
+        with open(os.path.join(combined_directory, "domainDict"), "r", encoding="utf-8") as file:
             domainDict = file.read()
     except FileNotFoundError:
         logger.error("文件不存在")
@@ -37,7 +37,7 @@ def get_dirDict():
     domainDict = ""
     try:
         # 尝试打开文件并读取内容
-        with open(os.path.join(combined_directory, "dirDict"), "r") as file:
+        with open(os.path.join(combined_directory, "dirDict"), "r", encoding="utf-8") as file:
             domainDict = file.read()
     except FileNotFoundError:
         logger.error("文件不存在")

core/util.py

@@ -403,7 +403,7 @@ async def get_search_query(name, request_data):
     if query == "" or query is None:
         return ""
     query = query[0]
-    filter_key = ['color', 'status', 'level', 'type', 'project']
+    filter_key = {'app':'app','color': 'color', 'status': 'status', 'level': 'level', 'type': 'type', 'project': 'project', 'port': 'port', 'protocol': ['protocol', 'type'], 'icon': 'faviconmmh3'}
     filter = request_data.get("filter", {})
     if filter:
         query["$and"] = []
@@ -412,7 +412,17 @@ async def get_search_query(name, request_data):
                 tmp_or = []
                 for v in filter[f]:
                     if v != "":
-                        tmp_or.append({f: v})
+                        if f == 'app':
+                            for ap_key in APP:
+                                if v == APP[ap_key]:
+                                    tmp_or.append({'webfinger': ap_key})
+                            tmp_or.append({'technologies': v})
+                        else:
+                            if type(filter_key[f]) is list:
+                                for li in filter_key[f]:
+                                    tmp_or.append({li: v})
+                            else:
+                                tmp_or.append({filter_key[f]: v})
                 if len(tmp_or) != 0:
                     query["$and"].append({"$or": tmp_or})
     if "$and" in query:

main.py

@@ -3,10 +3,13 @@ import time
 
 from loguru import logger
 import uvicorn
+from motor.motor_asyncio import AsyncIOMotorGridFSBucket
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.staticfiles import StaticFiles
 
 from core.config import *
+from core.default import get_dirDict, get_domainDict
+
 set_config()
 
 from core.db import get_mongo_db
@@ -32,6 +35,7 @@ from core.apscheduler_handler import scheduler
 
 async def update():
     async for db in get_mongo_db():
+        # 默认项目有个root_domain为空导致匹配上所有资产
         cursor = db.project.find({"root_domains": ""}, {"_id": 1, "root_domains": 1})
         async for document in cursor:
             logger.info("Update found empty root_domains")
@@ -45,6 +49,29 @@ async def update():
                 }
             }
             await db.project.update_one({"_id": document['_id']}, update_document)
+        # 修改目录字典存储方式
+        fs = AsyncIOMotorGridFSBucket(db)
+        result = await db.config.find_one({"name": "DirDic"})
+        if result:
+            await db.config.delete_one({"name": "DirDic"})
+            content = get_dirDict()
+            if content:
+                byte_content = content.encode('utf-8')
+                await fs.upload_from_stream('dirdict', byte_content)
+                logger.info("Document DirDict uploaded to GridFS.")
+            else:
+                logger.error("No dirdict content to upload.")
+        # 修改子域名字典存储方式
+        result = await db.config.find_one({"name": "DomainDic"})
+        if result:
+            await db.config.delete_one({"name": "DomainDic"})
+            content = get_domainDict()
+            if content:
+                byte_content = content.encode('utf-8')
+                await fs.upload_from_stream('DomainDic', byte_content)
+                logger.info("Document DomainDic uploaded to GridFS.")
+            else:
+                logger.error("No DomainDic content to upload.")
 
 
 @app.on_event("startup")