267 lines
9.9 KiB
Python
267 lines
9.9 KiB
Python
# -*- coding:utf-8 -*-
|
||
# @name: dictionary
|
||
# @auth: rainy-autumn@outlook.com
|
||
# @version:
|
||
from bson import ObjectId
|
||
from fastapi import APIRouter, Depends, File, UploadFile
|
||
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.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:
|
||
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))
|
||
|
||
|
||
@router.post("/subdomain/save")
|
||
async def save_subdomain_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': '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:
|
||
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")
|
||
async def get_port_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||
try:
|
||
page_index = request_data.get("pageIndex", 1)
|
||
page_size = request_data.get("pageSize", 10)
|
||
search = request_data.get("search", None) # 获取search参数
|
||
# Construct the search query
|
||
search_query = {}
|
||
if search:
|
||
search_regex = {"$regex": search, "$options": "i"} # Case-insensitive regex
|
||
search_query = {"$or": [{"name": search_regex}, {"value": search_regex}]}
|
||
total_count = await db.PortDict.count_documents(search_query)
|
||
|
||
# Perform pagination query
|
||
cursor: AsyncIOMotorCursor = db.PortDict.find(search_query).skip((page_index - 1) * page_size).limit(page_size)
|
||
result = await cursor.to_list(length=None)
|
||
|
||
# Process the result as needed
|
||
response_data = [{"id": str(doc["_id"]),"name": doc["name"], "value": doc["value"]} for doc in result]
|
||
return {
|
||
"code": 200,
|
||
"data": {
|
||
'list': response_data,
|
||
'total': total_count
|
||
}
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(str(e))
|
||
# Handle exceptions as needed
|
||
return {"message": "error","code":500}
|
||
|
||
@router.post("/port/upgrade")
|
||
async def upgrade_port_dict(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||
try:
|
||
# Extract values from request data
|
||
port_id = request_data.get("id")
|
||
name = request_data.get("name")
|
||
value = request_data.get("value")
|
||
|
||
# Update query based on rule_id
|
||
update_query = {"_id": ObjectId(port_id)}
|
||
|
||
# Values to be updated
|
||
update_values = {"$set": {"name": name, "value": value}}
|
||
|
||
# Perform the update
|
||
result = await db.PortDict.update_one(update_query, update_values)
|
||
await refresh_config('all', 'port')
|
||
if result:
|
||
return {"code": 200, "message": "SensitiveRule updated successfully"}
|
||
else:
|
||
return {"code": 404, "message": "SensitiveRule not found"}
|
||
|
||
except Exception as e:
|
||
logger.error(str(e))
|
||
# Handle exceptions as needed
|
||
return {"message": "error", "code": 500}
|
||
|
||
@router.post("/port/add")
|
||
async def add_port_dict(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||
try:
|
||
# Extract values from request data
|
||
name = request_data.get("name")
|
||
value = request_data.get("value",'')
|
||
if value == '':
|
||
return {"code": 400, "message": "value is null"}
|
||
# Create a new SensitiveRule document
|
||
new_port_dict = {
|
||
"name": name,
|
||
"value": value
|
||
}
|
||
|
||
# Insert the new document into the SensitiveRule collection
|
||
result = await db.PortDict.insert_one(new_port_dict)
|
||
await refresh_config('all', 'port')
|
||
# Check if the insertion was successful
|
||
if result.inserted_id:
|
||
return {"code": 200, "message": "Port dict added successfully"}
|
||
else:
|
||
return {"code": 400, "message": "Failed to add port dict"}
|
||
|
||
except Exception as e:
|
||
logger.error(str(e))
|
||
# Handle exceptions as needed
|
||
return {"message": "error", "code": 500}
|
||
|
||
@router.post("/port/delete")
|
||
async def delete_port_dict(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
|
||
try:
|
||
# Extract the list of IDs from the request_data dictionary
|
||
port_dict_ids = request_data.get("ids", [])
|
||
|
||
# Convert the provided rule_ids to ObjectId
|
||
obj_ids = [ObjectId(port_dict_id) for port_dict_id in port_dict_ids]
|
||
|
||
# Delete the SensitiveRule documents based on the provided IDs
|
||
result = await db.PortDict.delete_many({"_id": {"$in": obj_ids}})
|
||
await refresh_config('all', 'port')
|
||
# Check if the deletion was successful
|
||
if result.deleted_count > 0:
|
||
return {"code": 200, "message": "Port dict deleted successfully"}
|
||
else:
|
||
return {"code": 404, "message": "Port dict not found"}
|
||
|
||
except Exception as e:
|
||
logger.error(str(e))
|
||
# Handle exceptions as needed
|
||
return {"message": "error", "code": 500} |