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

add update、fix bug

This commit is contained in:
Autumn.home 2024-07-16 23:30:31 +08:00
parent afeffa6d3f
commit 689debfe87
8 changed files with 297 additions and 222 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
api/asset_info.py
View File

@ -55,7 +55,7 @@ async def asset_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Dep
if query == "":
return {"message": "Search condition parsing error", "code": 500}
total_count = await db['asset'].count_documents(query)
cursor: AsyncIOMotorCursor = ((db['asset'].find(query, {"_id": 0,
cursor = db['asset'].find(query, {"_id": 0,
"id": {"$toString": "$_id"},
"host": 1,
"url": 1,
@ -71,11 +71,8 @@ async def asset_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Dep
"raw": 1,
"timestamp": 1,
"iconcontent": 1
})
.skip((page_index - 1) * page_size)
.limit(page_size))
.sort([("timestamp", DESCENDING)]))
result = cursor.to_list(length=None)
}).skip((page_index - 1) * page_size).limit(page_size).sort([("timestamp", DESCENDING)])
result = await cursor.to_list(length=None)
result_list = []
for r in result:
tmp = {}

21
api/export.py
View File

@ -14,7 +14,7 @@ from starlette.responses import FileResponse
from api.users import verify_token
from motor.motor_asyncio import AsyncIOMotorCursor
from core.db import get_mongo_db
from core.db import get_mongo_db, get_project
import pandas as pd
from core.util import *
from pymongo import ASCENDING, DESCENDING, results
@ -151,7 +151,7 @@ async def fetch_data(db, collection, query, quantity, project_list):
{"$project": {"_id": 0, "vulnid": 0}}
]
cursor = await db[collection].aggregate(pipeline)
cursor = db[collection].aggregate(pipeline)
return cursor
@ -172,10 +172,21 @@ def flatten_dict(d):
items.append((k, v))
return dict(items)
def clean_string(value):
if isinstance(value, str):
# 过滤掉非法字符（ASCII码 < 32 或 >= 127）
return ''.join(char for char in value if 32 <= ord(char) < 127)
return value
async def export_data_from_mongodb(quantity, query, file_name, index):
logger.info("导出开始")
async for db in get_mongo_db():
try:
global Project_List
if len(Project_List) == 0:
await get_project(db)
cursor = await fetch_data(db, index, query, quantity, Project_List)
result = await cursor.to_list(length=None)
relative_path = f'file/{file_name}.xlsx'
@ -235,10 +246,10 @@ async def export_data_from_mongodb(quantity, query, file_name, index):
for doc in result:
flattened_doc = flatten_dict(doc)
if doc["type"] == "other":
row = [flattened_doc.get(col, "") for col in other_columns.keys()]
row = [clean_string(flattened_doc.get(col, "")) for col in other_columns.keys()]
other_ws.append(row)
else:
row = [flattened_doc.get(col, "") for col in http_columns.keys()]
row = [clean_string(flattened_doc.get(col, "")) for col in http_columns.keys()]
http_ws.append(row)
else:
columns = {}
@ -283,7 +294,7 @@ async def export_data_from_mongodb(quantity, query, file_name, index):
for doc in result:
flattened_doc = flatten_dict(doc)
row = [flattened_doc.get(col, "") for col in columns.keys()]
row = [clean_string(flattened_doc.get(col, "")) for col in columns.keys()]
ws.append(row)
try:
wb.save(file_path)

2
api/project.py
View File

@ -116,7 +116,7 @@ async def get_projects_all(db=Depends(get_mongo_db), _: dict = Depends(verify_to
async def update_project_count():
db = await get_mongo_db()
async for db in get_mongo_db():
cursor = db.project.find({}, {"_id": 0, "id": {"$toString": "$_id"}})
results = await cursor.to_list(length=None)

7
api/sensitive.py
View File

@ -256,7 +256,8 @@ async def get_sensitive_result_data2(request_data: dict, db=Depends(get_mongo_db
"time": 1,
"sid": 1,
"match": 1,
"color": 1
"color": 1,
"md5": 1
}
},
{
@ -267,7 +268,7 @@ async def get_sensitive_result_data2(request_data: dict, db=Depends(get_mongo_db
"_id": "$url",
"time": {"$first": "$time"}, # 记录相同url下最早插入数据的时间
"url": {"$first": "$url"},
"body_id": {"$last": {"$toString": "$_id"}}, # 记录相同url下最早插入数据的_id
"body_id": {"$last": {"$toString": "$md5"}}, # 记录相同url下最早插入数据的_id
"children": {
"$push": {
"id": {"$toString": "$_id"},
@ -316,7 +317,7 @@ async def get_sensitive_result_body_rules(request_data: dict, db=Depends(get_mon
return {"message": "ID is missing in the request data", "code": 400}
# Query the database for content based on ID
query = {"_id": ObjectId(sensitive_result_id)}
query = {"md5": sensitive_result_id}
doc = await db.SensitiveResult.find_one(query)
if not doc:

2
core/config.py
View File

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

2
core/db.py
View File

@ -94,7 +94,7 @@ async def create_database():
# {"name": "DomainDic", 'value': domainDict, 'type': 'domainDict'})
sensitive_data = get_sensitive()
collection = client[DATABASE_NAME]["SensitiveRule"]
if sensitiveList:
if sensitive_data:
await collection.insert_many(sensitive_data)
collection = client[DATABASE_NAME]["ScheduledTasks"]

327
dicts/ScopeSentry.SensitiveRule.json
View File

File diff suppressed because it is too large Load Diff

23
main.py
View File

@ -8,7 +8,7 @@ from starlette.middleware.base import BaseHTTPMiddleware
from starlette.staticfiles import StaticFiles
from core.config import *
from core.default import get_dirDict, get_domainDict
from core.default import get_dirDict, get_domainDict, get_sensitive
set_config()
@ -35,6 +35,19 @@ from core.apscheduler_handler import scheduler
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"]
update = result["update"]
if version < float(VERSION):
update = False
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:
@ -73,6 +86,14 @@ async def update():
else:
logger.error("No DomainDic content to upload.")
# 更新敏感信息
await db.SensitiveRule.delete_many({})
sensitive_data = get_sensitive()
collection = db["SensitiveRule"]
if sensitive_data:
await collection.insert_many(sensitive_data)
await db.config.update_one({"name": "version"}, {"$set": {"update": True, "version": float(VERSION)}})
@app.on_event("startup")
async def startup_db_client():