From 30ed71b838c5a47b4b9b0858a40a8cdcc7478222 Mon Sep 17 00:00:00 2001 From: "Autumn.home" Date: Sun, 7 Jul 2024 23:21:34 +0800 Subject: [PATCH] add --- api/dictionary.py | 155 ++++++++++++++++++++++++++++++++++------------ core/db.py | 12 ++-- core/default.py | 4 +- core/util.py | 14 ++++- main.py | 27 ++++++++ requirements.txt | Bin 1830 -> 1866 bytes 6 files changed, 164 insertions(+), 48 deletions(-) diff --git a/api/dictionary.py b/api/dictionary.py index 3e1a159..ddbdbf9 100644 --- a/api/dictionary.py +++ b/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" - result = await db.config.find_one({"name": "DomainDic"}) - return { - "code": 200, - "data": { - "dict": result.get("value", '') - } - } + fs = AsyncIOMotorGridFSBucket(db) + # 查找文件 + 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" - 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"} + content = await file.read() + fs = AsyncIOMotorGridFSBucket(db) + 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" - result = await db.config.find_one({"name": "DirDic"}) - return { - "code": 200, - "data": { - "dict": result.get("value", '') - } - } + fs = AsyncIOMotorGridFSBucket(db) + # 查找文件 + 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" - 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"} + content = await file.read() + fs = AsyncIOMotorGridFSBucket(db) + 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 diff --git a/core/db.py b/core/db.py index 5df1c44..f825112 100644 --- a/core/db.py +++ b/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() - await collection.insert_one( - {"name": "DirDic", 'value': dirDict, 'type': 'dirDict'}) + # dirDict = get_dirDict() + # await collection.insert_one( + # {"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() - await collection.insert_one( - {"name": "DomainDic", 'value': domainDict, 'type': 'domainDict'}) + # domainDict = get_domainDict() + # await collection.insert_one( + # {"name": "DomainDic", 'value': domainDict, 'type': 'domainDict'}) sensitive_data = get_sensitive() collection = client[DATABASE_NAME]["SensitiveRule"] if sensitiveList: diff --git a/core/default.py b/core/default.py index 092aaf3..acd1fef 100644 --- a/core/default.py +++ b/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("文件不存在") diff --git a/core/util.py b/core/util.py index 2ecca04..f92bd6f 100644 --- a/core/util.py +++ b/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: diff --git a/main.py b/main.py index ed0973d..2eacdb4 100644 --- a/main.py +++ b/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") diff --git a/requirements.txt b/requirements.txt index 70a1b73364eabc1d1baea026ac15c091a3587862..7a5899f4f537d04da1a3243e749bb68dc14ea3d0 100644 GIT binary patch delta 42 wcmZ3+cZzR=>Px# delta 10 RcmX@bw~TLs+(su3b^sNX13dr$