liangsheng8708
/
ScopeSentry
mirror of https://mirror.ghproxy.com/https://github.com/Autumn-27/ScopeSentry.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

修改字典相关逻辑

This commit is contained in:
Autumn.home 2024-09-08 21:32:34 +08:00
parent f12ebc25e4
commit e2f8f97c5d
10 changed files with 373 additions and 147 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
api/configuration.py
View File

@ -19,6 +19,7 @@ from core.util import generate_random_string
router = APIRouter()
@router.get("/subfinder/data")
async def get_subfinder_data(db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
try:

88
api/dictionary.py
View File

@ -67,94 +67,7 @@ async def save_subdomain_data(file: UploadFile = File(...), db=Depends(get_mongo
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)):
@ -243,6 +156,7 @@ async def add_port_dict(request_data: dict, db=Depends(get_mongo_db), _: dict =
# 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:

14
api/dictionary/__init__.py Normal file
View File

@ -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")

123
api/dictionary/dir.py Normal file
View File

@ -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"}

123
api/dictionary/subdomain.py Normal file
View File

@ -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"}

6
core/apscheduler_handler.py
View File

@ -11,9 +11,9 @@ from core.config import *
mongo_config = {
'host': MONGODB_IP,
'port': int(MONGODB_PORT),
'username': DATABASE_USER,
'password': DATABASE_PASSWORD,
'database': DATABASE_NAME,
'username': str(DATABASE_USER),
'password': str(DATABASE_PASSWORD),
'database': str(DATABASE_NAME),
'collection': 'apscheduler'
}
jobstores = {

2
core/config.py
View File

@ -8,7 +8,7 @@ import string
import yaml
VERSION = "1.4"
VERSION = "1.5"
UPDATEURL = "http://update.scope-sentry.top"
REMOTE_REPO_URL = "https://github.com/Autumn-27/ScopeSentry.git"
SECRET_KEY = "ScopeSentry-15847412364125411"

9
core/db.py
View File

@ -53,8 +53,9 @@ async def create_database():
collection = db["config"]
# 扫描模块配置
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配置
collection = db["config"]
# 插入一条数据
@ -70,12 +71,12 @@ async def create_database():
content = get_dirDict()
if content:
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()
if content:
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.")
await collection.insert_one(

87
core/update.py
View File

@ -4,3 +4,90 @@
# @contact : rainy-autumn@outlook.com
# @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
View File

@ -9,6 +9,7 @@ from starlette.staticfiles import StaticFiles
from core.config import *
from core.default import get_dirDict, get_domainDict, get_sensitive
from core.update import update14, update15
set_config()
@ -37,7 +38,6 @@ async def update():
async for db in get_mongo_db():
# 判断版本
result = await db.config.find_one({"name": "version"})
version = 0
update = False
if result is not None:
version = result["version"]
@ -47,56 +47,11 @@ async def update():
else:
await db.config.insert_one({"name": "version", "version": float(VERSION), "update": False})
version = float(VERSION)
if version <= 1.4 and update is False:
# 默认项目有个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)}})
if update is False:
if version < 1.4:
await update14(db)
if version < 1.5:
await update15(db)
@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
pat, flag = await get_page_monitoring_time()
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())
@ -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, \
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(sensitive.router, prefix='/api')
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(configuration.router, prefix='/api/configuration')
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(project_aggregation.router, prefix='/api/project_aggregation')
app.mount("/assets", StaticFiles(directory="static/assets"), name="assets")
@app.get("/logo.png", response_class=FileResponse)
async def get_logo(request: Request):
return FileResponse("static/logo.png")
@ -160,6 +121,8 @@ async def get_logo(request: Request):
@app.get("/favicon.ico", response_class=FileResponse)
async def get_favicon(request: Request):
return FileResponse("static/favicon.ico")
# @app.middleware("http")
# async def process_http_requests(request, call_next):
# url = str(request.url)