修改字典相关逻辑
This commit is contained in:
parent
f12ebc25e4
commit
e2f8f97c5d
|
@ -19,6 +19,7 @@ from core.util import generate_random_string
|
||||||
|
|
||||||
router = APIRouter()
|
router = APIRouter()
|
||||||
|
|
||||||
|
|
||||||
@router.get("/subfinder/data")
|
@router.get("/subfinder/data")
|
||||||
async def get_subfinder_data(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
async def get_subfinder_data(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
try:
|
try:
|
||||||
|
|
|
@ -67,94 +67,7 @@ async def save_subdomain_data(file: UploadFile = File(...), db=Depends(get_mongo
|
||||||
logger.error(str(e))
|
logger.error(str(e))
|
||||||
# Handle exceptions as needed
|
# Handle exceptions as needed
|
||||||
return {"message": "error", "code": 500}
|
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:
|
|
||||||
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))
|
|
||||||
|
|
||||||
# @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_dir_data(file: UploadFile = File(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
|
||||||
try:
|
|
||||||
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
|
|
||||||
return {"message": "error", "code": 500}
|
|
||||||
|
|
||||||
@router.post("/port/data")
|
@router.post("/port/data")
|
||||||
async def get_port_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
async def get_port_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
@ -243,6 +156,7 @@ async def add_port_dict(request_data: dict, db=Depends(get_mongo_db), _: dict =
|
||||||
# Handle exceptions as needed
|
# Handle exceptions as needed
|
||||||
return {"message": "error", "code": 500}
|
return {"message": "error", "code": 500}
|
||||||
|
|
||||||
|
|
||||||
@router.post("/port/delete")
|
@router.post("/port/delete")
|
||||||
async def delete_port_dict(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
async def delete_port_dict(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
try:
|
try:
|
||||||
|
|
|
@ -0,0 +1,14 @@
|
||||||
|
# -------------------------------------
|
||||||
|
# @file : __init__.py
|
||||||
|
# @author : Autumn
|
||||||
|
# @contact : rainy-autumn@outlook.com
|
||||||
|
# @time : 2024/9/8 15:10
|
||||||
|
# -------------------------------------------
|
||||||
|
from fastapi import APIRouter
|
||||||
|
|
||||||
|
from .dir import router as dir_router
|
||||||
|
from .subdomain import router as subdoamin_router
|
||||||
|
router = APIRouter()
|
||||||
|
|
||||||
|
router.include_router(dir_router, prefix="/dir")
|
||||||
|
router.include_router(subdoamin_router, prefix="/subdomain")
|
|
@ -0,0 +1,123 @@
|
||||||
|
# -------------------------------------
|
||||||
|
# @file : dir.py
|
||||||
|
# @author : Autumn
|
||||||
|
# @contact : rainy-autumn@outlook.com
|
||||||
|
# @time : 2024/9/8 15:10
|
||||||
|
# -------------------------------------------
|
||||||
|
import re
|
||||||
|
|
||||||
|
from bson import ObjectId
|
||||||
|
from fastapi import APIRouter, Depends, File, UploadFile, Query, Form
|
||||||
|
from starlette.responses import StreamingResponse
|
||||||
|
|
||||||
|
from api.users import verify_token
|
||||||
|
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.post("/list")
|
||||||
|
async def list(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
cursor = fs.find({"filename": {"$regex": f"^dir_"}})
|
||||||
|
files_data = await cursor.to_list(length=None)
|
||||||
|
result_list = []
|
||||||
|
for file in files_data:
|
||||||
|
result_list.append({
|
||||||
|
"id": str(file["_id"]),
|
||||||
|
"name": file["filename"].replace("dir_", ""),
|
||||||
|
"size": "{:.2f}".format(file["length"] / (1024 * 1024)),
|
||||||
|
})
|
||||||
|
return {
|
||||||
|
"code": 200,
|
||||||
|
"data": {
|
||||||
|
'list': result_list,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/save")
|
||||||
|
async def save_dir_data(file: UploadFile = File(...), id: str = Query(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
try:
|
||||||
|
content = await file.read()
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
|
||||||
|
# 查找文件,根据传入的 id
|
||||||
|
old_file = await fs.find({'_id': ObjectId(id)}).to_list(1)
|
||||||
|
|
||||||
|
if old_file:
|
||||||
|
# 如果找到文件,删除旧文件
|
||||||
|
await fs.delete(ObjectId(id))
|
||||||
|
print(f"File with id {id} deleted.")
|
||||||
|
|
||||||
|
# 将新文件内容存储到 GridFS 中,并保存新的文件 ID
|
||||||
|
await fs.upload_from_stream(
|
||||||
|
old_file[0]['filename'], # 使用原文件名存储
|
||||||
|
content # 文件内容
|
||||||
|
)
|
||||||
|
return {"code": 200, "message": "upload successful"}
|
||||||
|
else:
|
||||||
|
return {"message": "not found", "code": 500}
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(str(e))
|
||||||
|
# Handle exceptions as needed
|
||||||
|
return {"message": "error", "code": 500}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/add")
|
||||||
|
async def add_dir_data(file: UploadFile = File(...), name: str = Form(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
try:
|
||||||
|
if not re.match(r'^[a-zA-Z0-9_-]+$', name):
|
||||||
|
return {"message": "Name must contain only letters", "code": 500}
|
||||||
|
content = await file.read()
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
await fs.upload_from_stream(
|
||||||
|
f"dir_{name}", # 使用原文件名存储
|
||||||
|
content # 文件内容
|
||||||
|
)
|
||||||
|
return {"code": 200, "message": "upload successful"}
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(str(e))
|
||||||
|
# Handle exceptions as needed
|
||||||
|
return {"message": "error", "code": 500}
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/download")
|
||||||
|
async def get_dir_data(id: str = Query(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
try:
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
|
||||||
|
# 查找文件
|
||||||
|
file_doc = await fs.find({'_id': ObjectId(id)}).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)
|
||||||
|
filename = file_doc[0]['filename'].replace("dir_", "")
|
||||||
|
# 返回文件流
|
||||||
|
return StreamingResponse(grid_out, media_type="application/octet-stream",
|
||||||
|
headers={"Content-Disposition": f"attachment; filename={filename}"})
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(str(e))
|
||||||
|
return {"message": "error", "code": 500}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/delete")
|
||||||
|
async def delete_dir_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
dir_dict_ids = request_data.get("ids", [])
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
for id in dir_dict_ids:
|
||||||
|
try:
|
||||||
|
file_id = ObjectId(id)
|
||||||
|
# 删除文件
|
||||||
|
await fs.delete(file_id)
|
||||||
|
print(f"文件 {id} 已删除")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"删除文件 {id} 时发生错误: {e}")
|
||||||
|
logger.error(str(e))
|
||||||
|
return {"message": "error", "code": 500}
|
||||||
|
return {"code": 200, "message": "delete file successful"}
|
|
@ -0,0 +1,123 @@
|
||||||
|
# -------------------------------------
|
||||||
|
# @file : subdomain.py
|
||||||
|
# @author : Autumn
|
||||||
|
# @contact : rainy-autumn@outlook.com
|
||||||
|
# @time : 2024/9/8 21:19
|
||||||
|
# -------------------------------------------
|
||||||
|
import re
|
||||||
|
|
||||||
|
from bson import ObjectId
|
||||||
|
from fastapi import APIRouter, Depends, File, UploadFile, Query, Form
|
||||||
|
from starlette.responses import StreamingResponse
|
||||||
|
|
||||||
|
from api.users import verify_token
|
||||||
|
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.post("/list")
|
||||||
|
async def list(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
cursor = fs.find({"filename": {"$regex": f"^domain_"}})
|
||||||
|
files_data = await cursor.to_list(length=None)
|
||||||
|
result_list = []
|
||||||
|
for file in files_data:
|
||||||
|
result_list.append({
|
||||||
|
"id": str(file["_id"]),
|
||||||
|
"name": file["filename"].replace("domain_", ""),
|
||||||
|
"size": "{:.2f}".format(file["length"] / (1024 * 1024)),
|
||||||
|
})
|
||||||
|
return {
|
||||||
|
"code": 200,
|
||||||
|
"data": {
|
||||||
|
'list': result_list,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/save")
|
||||||
|
async def save_dir_data(file: UploadFile = File(...), id: str = Query(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
try:
|
||||||
|
content = await file.read()
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
|
||||||
|
# 查找文件,根据传入的 id
|
||||||
|
old_file = await fs.find({'_id': ObjectId(id)}).to_list(1)
|
||||||
|
|
||||||
|
if old_file:
|
||||||
|
# 如果找到文件,删除旧文件
|
||||||
|
await fs.delete(ObjectId(id))
|
||||||
|
print(f"File with id {id} deleted.")
|
||||||
|
|
||||||
|
# 将新文件内容存储到 GridFS 中,并保存新的文件 ID
|
||||||
|
await fs.upload_from_stream(
|
||||||
|
old_file[0]['filename'], # 使用原文件名存储
|
||||||
|
content # 文件内容
|
||||||
|
)
|
||||||
|
return {"code": 200, "message": "upload successful"}
|
||||||
|
else:
|
||||||
|
return {"message": "not found", "code": 500}
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(str(e))
|
||||||
|
# Handle exceptions as needed
|
||||||
|
return {"message": "error", "code": 500}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/add")
|
||||||
|
async def add_subdomain_data(file: UploadFile = File(...), name: str = Form(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
try:
|
||||||
|
if not re.match(r'^[a-zA-Z0-9_-]+$', name):
|
||||||
|
return {"message": "Name must contain only letters", "code": 500}
|
||||||
|
content = await file.read()
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
await fs.upload_from_stream(
|
||||||
|
f"domain_{name}", # 使用原文件名存储
|
||||||
|
content # 文件内容
|
||||||
|
)
|
||||||
|
return {"code": 200, "message": "upload successful"}
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(str(e))
|
||||||
|
# Handle exceptions as needed
|
||||||
|
return {"message": "error", "code": 500}
|
||||||
|
|
||||||
|
|
||||||
|
@router.get("/download")
|
||||||
|
async def get_subdomain_data(id: str = Query(...), db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
try:
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
|
||||||
|
# 查找文件
|
||||||
|
file_doc = await fs.find({'_id': ObjectId(id)}).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)
|
||||||
|
filename = file_doc[0]['filename'].replace("domain_", "")
|
||||||
|
# 返回文件流
|
||||||
|
return StreamingResponse(grid_out, media_type="application/octet-stream",
|
||||||
|
headers={"Content-Disposition": f"attachment; filename={filename}"})
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(str(e))
|
||||||
|
return {"message": "error", "code": 500}
|
||||||
|
|
||||||
|
|
||||||
|
@router.post("/delete")
|
||||||
|
async def delete_subdomain_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||||||
|
dir_dict_ids = request_data.get("ids", [])
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
for id in dir_dict_ids:
|
||||||
|
try:
|
||||||
|
file_id = ObjectId(id)
|
||||||
|
# 删除文件
|
||||||
|
await fs.delete(file_id)
|
||||||
|
print(f"文件 {id} 已删除")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"删除文件 {id} 时发生错误: {e}")
|
||||||
|
logger.error(str(e))
|
||||||
|
return {"message": "error", "code": 500}
|
||||||
|
return {"code": 200, "message": "delete file successful"}
|
|
@ -11,9 +11,9 @@ from core.config import *
|
||||||
mongo_config = {
|
mongo_config = {
|
||||||
'host': MONGODB_IP,
|
'host': MONGODB_IP,
|
||||||
'port': int(MONGODB_PORT),
|
'port': int(MONGODB_PORT),
|
||||||
'username': DATABASE_USER,
|
'username': str(DATABASE_USER),
|
||||||
'password': DATABASE_PASSWORD,
|
'password': str(DATABASE_PASSWORD),
|
||||||
'database': DATABASE_NAME,
|
'database': str(DATABASE_NAME),
|
||||||
'collection': 'apscheduler'
|
'collection': 'apscheduler'
|
||||||
}
|
}
|
||||||
jobstores = {
|
jobstores = {
|
||||||
|
|
|
@ -8,7 +8,7 @@ import string
|
||||||
|
|
||||||
import yaml
|
import yaml
|
||||||
|
|
||||||
VERSION = "1.4"
|
VERSION = "1.5"
|
||||||
UPDATEURL = "http://update.scope-sentry.top"
|
UPDATEURL = "http://update.scope-sentry.top"
|
||||||
REMOTE_REPO_URL = "https://github.com/Autumn-27/ScopeSentry.git"
|
REMOTE_REPO_URL = "https://github.com/Autumn-27/ScopeSentry.git"
|
||||||
SECRET_KEY = "ScopeSentry-15847412364125411"
|
SECRET_KEY = "ScopeSentry-15847412364125411"
|
||||||
|
|
|
@ -53,8 +53,9 @@ async def create_database():
|
||||||
collection = db["config"]
|
collection = db["config"]
|
||||||
# 扫描模块配置
|
# 扫描模块配置
|
||||||
await collection.insert_one(
|
await collection.insert_one(
|
||||||
{"name": "Modules", 'value': ModulesConfig, 'type': 'system'})
|
{"name": "ModulesConfig", 'value': ModulesConfig, 'type': 'system'})
|
||||||
|
await collection.insert_one(
|
||||||
|
{"name": "timezone", 'value': 'Asia/Shanghai', 'type': 'system'})
|
||||||
# subfinder配置
|
# subfinder配置
|
||||||
collection = db["config"]
|
collection = db["config"]
|
||||||
# 插入一条数据
|
# 插入一条数据
|
||||||
|
@ -70,12 +71,12 @@ async def create_database():
|
||||||
content = get_dirDict()
|
content = get_dirDict()
|
||||||
if content:
|
if content:
|
||||||
byte_content = content.encode('utf-8')
|
byte_content = content.encode('utf-8')
|
||||||
await fs.upload_from_stream('dirdict', byte_content)
|
await fs.upload_from_stream('dir_default', byte_content)
|
||||||
# 子域名字典
|
# 子域名字典
|
||||||
content = get_domainDict()
|
content = get_domainDict()
|
||||||
if content:
|
if content:
|
||||||
byte_content = content.encode('utf-8')
|
byte_content = content.encode('utf-8')
|
||||||
await fs.upload_from_stream('DomainDic', byte_content)
|
await fs.upload_from_stream('domain_default', byte_content)
|
||||||
logger.info("Document DomainDic uploaded to GridFS.")
|
logger.info("Document DomainDic uploaded to GridFS.")
|
||||||
|
|
||||||
await collection.insert_one(
|
await collection.insert_one(
|
||||||
|
|
|
@ -4,3 +4,90 @@
|
||||||
# @contact : rainy-autumn@outlook.com
|
# @contact : rainy-autumn@outlook.com
|
||||||
# @time : 2024/6/16 14:14
|
# @time : 2024/6/16 14:14
|
||||||
# -------------------------------------------
|
# -------------------------------------------
|
||||||
|
from loguru import logger
|
||||||
|
from motor.motor_asyncio import AsyncIOMotorGridFSBucket
|
||||||
|
|
||||||
|
from core.config import VERSION
|
||||||
|
from core.default import get_dirDict, get_domainDict, get_sensitive, ModulesConfig
|
||||||
|
|
||||||
|
|
||||||
|
async def update14(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")
|
||||||
|
root_domain = []
|
||||||
|
for root in document["root_domains"]:
|
||||||
|
if root != "":
|
||||||
|
root_domain.append(root)
|
||||||
|
update_document = {
|
||||||
|
"$set": {
|
||||||
|
"root_domains": root_domain,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
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.")
|
||||||
|
|
||||||
|
# 更新敏感信息
|
||||||
|
sensitive_data = get_sensitive()
|
||||||
|
collection = db["SensitiveRule"]
|
||||||
|
if sensitive_data:
|
||||||
|
for s in sensitive_data:
|
||||||
|
await collection.update_one(
|
||||||
|
{"name": s['name']},
|
||||||
|
{"$set": s},
|
||||||
|
upsert=True
|
||||||
|
)
|
||||||
|
await db.config.update_one({"name": "version"}, {"$set": {"update": True, "version": float(VERSION)}})
|
||||||
|
|
||||||
|
|
||||||
|
async def update15(db):
|
||||||
|
await db.config.insert_one(
|
||||||
|
{"name": "ModulesConfig", 'value': ModulesConfig, 'type': 'system'})
|
||||||
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
|
# 更新目录扫描字典
|
||||||
|
file_docs = await fs.files.find({"filename": "dirdict"}).to_list(1)
|
||||||
|
if not file_docs:
|
||||||
|
logger.error("File dirdict not found")
|
||||||
|
return
|
||||||
|
file_doc = file_docs[0]
|
||||||
|
new_filename = "dir_default"
|
||||||
|
# 更新文件名
|
||||||
|
await fs.files.update_one(
|
||||||
|
{"_id": file_doc["_id"]},
|
||||||
|
{"$set": {"filename": new_filename}}
|
||||||
|
)
|
||||||
|
|
||||||
|
# 更新子域名字典
|
||||||
|
file_docs = await fs.files.find({"filename": "DomainDic"}).to_list(1)
|
||||||
|
if not file_docs:
|
||||||
|
logger.error("File DomainDic not found")
|
||||||
|
file_doc = file_docs[0]
|
||||||
|
new_filename = "domain_default"
|
||||||
|
# 更新文件名
|
||||||
|
await fs.files.update_one(
|
||||||
|
{"_id": file_doc["_id"]},
|
||||||
|
{"$set": {"filename": new_filename}}
|
||||||
|
)
|
||||||
|
|
67
main.py
67
main.py
|
@ -9,6 +9,7 @@ from starlette.staticfiles import StaticFiles
|
||||||
|
|
||||||
from core.config import *
|
from core.config import *
|
||||||
from core.default import get_dirDict, get_domainDict, get_sensitive
|
from core.default import get_dirDict, get_domainDict, get_sensitive
|
||||||
|
from core.update import update14, update15
|
||||||
|
|
||||||
set_config()
|
set_config()
|
||||||
|
|
||||||
|
@ -37,7 +38,6 @@ async def update():
|
||||||
async for db in get_mongo_db():
|
async for db in get_mongo_db():
|
||||||
# 判断版本
|
# 判断版本
|
||||||
result = await db.config.find_one({"name": "version"})
|
result = await db.config.find_one({"name": "version"})
|
||||||
version = 0
|
|
||||||
update = False
|
update = False
|
||||||
if result is not None:
|
if result is not None:
|
||||||
version = result["version"]
|
version = result["version"]
|
||||||
|
@ -47,56 +47,11 @@ async def update():
|
||||||
else:
|
else:
|
||||||
await db.config.insert_one({"name": "version", "version": float(VERSION), "update": False})
|
await db.config.insert_one({"name": "version", "version": float(VERSION), "update": False})
|
||||||
version = float(VERSION)
|
version = float(VERSION)
|
||||||
if version <= 1.4 and update is False:
|
if update is False:
|
||||||
# 默认项目有个root_domain为空导致匹配上所有资产
|
if version < 1.4:
|
||||||
cursor = db.project.find({"root_domains": ""}, {"_id": 1, "root_domains": 1})
|
await update14(db)
|
||||||
async for document in cursor:
|
if version < 1.5:
|
||||||
logger.info("Update found empty root_domains")
|
await update15(db)
|
||||||
root_domain = []
|
|
||||||
for root in document["root_domains"]:
|
|
||||||
if root != "":
|
|
||||||
root_domain.append(root)
|
|
||||||
update_document = {
|
|
||||||
"$set": {
|
|
||||||
"root_domains": root_domain,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
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.")
|
|
||||||
|
|
||||||
# 更新敏感信息
|
|
||||||
sensitive_data = get_sensitive()
|
|
||||||
collection = db["SensitiveRule"]
|
|
||||||
if sensitive_data:
|
|
||||||
for s in sensitive_data:
|
|
||||||
await collection.update_one(
|
|
||||||
{"name": s['name']},
|
|
||||||
{"$set": s},
|
|
||||||
upsert=True
|
|
||||||
)
|
|
||||||
await db.config.update_one({"name": "version"}, {"$set": {"update": True, "version": float(VERSION)}})
|
|
||||||
|
|
||||||
|
|
||||||
@app.on_event("startup")
|
@app.on_event("startup")
|
||||||
|
@ -116,7 +71,8 @@ async def startup_db_client():
|
||||||
from api.scheduled_tasks import get_page_monitoring_time, create_page_monitoring_task
|
from api.scheduled_tasks import get_page_monitoring_time, create_page_monitoring_task
|
||||||
pat, flag = await get_page_monitoring_time()
|
pat, flag = await get_page_monitoring_time()
|
||||||
if flag:
|
if flag:
|
||||||
scheduler.add_job(create_page_monitoring_task, 'interval', hours=pat, id='page_monitoring', jobstore='mongo')
|
scheduler.add_job(create_page_monitoring_task, 'interval', hours=pat, id='page_monitoring',
|
||||||
|
jobstore='mongo')
|
||||||
asyncio.create_task(subscribe_log_channel())
|
asyncio.create_task(subscribe_log_channel())
|
||||||
|
|
||||||
|
|
||||||
|
@ -132,9 +88,12 @@ os.chdir(os.path.dirname(os.path.abspath(__file__)))
|
||||||
from api import users, sensitive, dictionary, poc, configuration, fingerprint, node, project, task, asset_info, \
|
from api import users, sensitive, dictionary, poc, configuration, fingerprint, node, project, task, asset_info, \
|
||||||
page_monitoring, vulnerability, SubdoaminTaker, scheduled_tasks, notification, system, export, project_aggregation
|
page_monitoring, vulnerability, SubdoaminTaker, scheduled_tasks, notification, system, export, project_aggregation
|
||||||
|
|
||||||
|
from api.dictionary import router as dictionary_router
|
||||||
|
|
||||||
app.include_router(users.router, prefix='/api')
|
app.include_router(users.router, prefix='/api')
|
||||||
app.include_router(sensitive.router, prefix='/api')
|
app.include_router(sensitive.router, prefix='/api')
|
||||||
app.include_router(dictionary.router, prefix='/api/dictionary')
|
app.include_router(dictionary.router, prefix='/api/dictionary')
|
||||||
|
app.include_router(dictionary_router, prefix='/api/dictionary')
|
||||||
app.include_router(poc.router, prefix='/api')
|
app.include_router(poc.router, prefix='/api')
|
||||||
app.include_router(configuration.router, prefix='/api/configuration')
|
app.include_router(configuration.router, prefix='/api/configuration')
|
||||||
app.include_router(fingerprint.router, prefix='/api')
|
app.include_router(fingerprint.router, prefix='/api')
|
||||||
|
@ -152,6 +111,8 @@ app.include_router(system.router, prefix='/api')
|
||||||
app.include_router(export.router, prefix='/api')
|
app.include_router(export.router, prefix='/api')
|
||||||
app.include_router(project_aggregation.router, prefix='/api/project_aggregation')
|
app.include_router(project_aggregation.router, prefix='/api/project_aggregation')
|
||||||
app.mount("/assets", StaticFiles(directory="static/assets"), name="assets")
|
app.mount("/assets", StaticFiles(directory="static/assets"), name="assets")
|
||||||
|
|
||||||
|
|
||||||
@app.get("/logo.png", response_class=FileResponse)
|
@app.get("/logo.png", response_class=FileResponse)
|
||||||
async def get_logo(request: Request):
|
async def get_logo(request: Request):
|
||||||
return FileResponse("static/logo.png")
|
return FileResponse("static/logo.png")
|
||||||
|
@ -160,6 +121,8 @@ async def get_logo(request: Request):
|
||||||
@app.get("/favicon.ico", response_class=FileResponse)
|
@app.get("/favicon.ico", response_class=FileResponse)
|
||||||
async def get_favicon(request: Request):
|
async def get_favicon(request: Request):
|
||||||
return FileResponse("static/favicon.ico")
|
return FileResponse("static/favicon.ico")
|
||||||
|
|
||||||
|
|
||||||
# @app.middleware("http")
|
# @app.middleware("http")
|
||||||
# async def process_http_requests(request, call_next):
|
# async def process_http_requests(request, call_next):
|
||||||
# url = str(request.url)
|
# url = str(request.url)
|
||||||
|
|
Loading…
Reference in New Issue