add
This commit is contained in:
parent
b7790f589d
commit
30ed71b838
|
@ -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
|
||||
|
|
12
core/db.py
12
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:
|
||||
|
|
|
@ -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("文件不存在")
|
||||
|
|
14
core/util.py
14
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:
|
||||
|
|
27
main.py
27
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")
|
||||
|
|
BIN
requirements.txt
BIN
requirements.txt
Binary file not shown.
Loading…
Reference in New Issue