add update、fix bug
This commit is contained in:
parent
afeffa6d3f
commit
689debfe87
|
@ -55,27 +55,24 @@ async def asset_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Dep
|
||||||
if query == "":
|
if query == "":
|
||||||
return {"message": "Search condition parsing error", "code": 500}
|
return {"message": "Search condition parsing error", "code": 500}
|
||||||
total_count = await db['asset'].count_documents(query)
|
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"},
|
"id": {"$toString": "$_id"},
|
||||||
"host": 1,
|
"host": 1,
|
||||||
"url": 1,
|
"url": 1,
|
||||||
"ip": 1,
|
"ip": 1,
|
||||||
"port": 1,
|
"port": 1,
|
||||||
"protocol": 1,
|
"protocol": 1,
|
||||||
"type": 1,
|
"type": 1,
|
||||||
"title": 1,
|
"title": 1,
|
||||||
"statuscode": 1,
|
"statuscode": 1,
|
||||||
"rawheaders": 1,
|
"rawheaders": 1,
|
||||||
"webfinger": 1,
|
"webfinger": 1,
|
||||||
"technologies": 1,
|
"technologies": 1,
|
||||||
"raw": 1,
|
"raw": 1,
|
||||||
"timestamp": 1,
|
"timestamp": 1,
|
||||||
"iconcontent": 1
|
"iconcontent": 1
|
||||||
})
|
}).skip((page_index - 1) * page_size).limit(page_size).sort([("timestamp", DESCENDING)])
|
||||||
.skip((page_index - 1) * page_size)
|
result = await cursor.to_list(length=None)
|
||||||
.limit(page_size))
|
|
||||||
.sort([("timestamp", DESCENDING)]))
|
|
||||||
result = cursor.to_list(length=None)
|
|
||||||
result_list = []
|
result_list = []
|
||||||
for r in result:
|
for r in result:
|
||||||
tmp = {}
|
tmp = {}
|
||||||
|
|
|
@ -14,7 +14,7 @@ from starlette.responses import FileResponse
|
||||||
|
|
||||||
from api.users import verify_token
|
from api.users import verify_token
|
||||||
from motor.motor_asyncio import AsyncIOMotorCursor
|
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
|
import pandas as pd
|
||||||
from core.util import *
|
from core.util import *
|
||||||
from pymongo import ASCENDING, DESCENDING, results
|
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}}
|
{"$project": {"_id": 0, "vulnid": 0}}
|
||||||
]
|
]
|
||||||
|
|
||||||
cursor = await db[collection].aggregate(pipeline)
|
cursor = db[collection].aggregate(pipeline)
|
||||||
return cursor
|
return cursor
|
||||||
|
|
||||||
|
|
||||||
|
@ -172,10 +172,21 @@ def flatten_dict(d):
|
||||||
items.append((k, v))
|
items.append((k, v))
|
||||||
return dict(items)
|
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):
|
async def export_data_from_mongodb(quantity, query, file_name, index):
|
||||||
logger.info("导出开始")
|
logger.info("导出开始")
|
||||||
async for db in get_mongo_db():
|
async for db in get_mongo_db():
|
||||||
try:
|
try:
|
||||||
|
global Project_List
|
||||||
|
if len(Project_List) == 0:
|
||||||
|
await get_project(db)
|
||||||
cursor = await fetch_data(db, index, query, quantity, Project_List)
|
cursor = await fetch_data(db, index, query, quantity, Project_List)
|
||||||
result = await cursor.to_list(length=None)
|
result = await cursor.to_list(length=None)
|
||||||
relative_path = f'file/{file_name}.xlsx'
|
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:
|
for doc in result:
|
||||||
flattened_doc = flatten_dict(doc)
|
flattened_doc = flatten_dict(doc)
|
||||||
if doc["type"] == "other":
|
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)
|
other_ws.append(row)
|
||||||
else:
|
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)
|
http_ws.append(row)
|
||||||
else:
|
else:
|
||||||
columns = {}
|
columns = {}
|
||||||
|
@ -283,7 +294,7 @@ async def export_data_from_mongodb(quantity, query, file_name, index):
|
||||||
|
|
||||||
for doc in result:
|
for doc in result:
|
||||||
flattened_doc = flatten_dict(doc)
|
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)
|
ws.append(row)
|
||||||
try:
|
try:
|
||||||
wb.save(file_path)
|
wb.save(file_path)
|
||||||
|
|
|
@ -116,23 +116,23 @@ async def get_projects_all(db=Depends(get_mongo_db), _: dict = Depends(verify_to
|
||||||
|
|
||||||
|
|
||||||
async def update_project_count():
|
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"}})
|
cursor = db.project.find({}, {"_id": 0, "id": {"$toString": "$_id"}})
|
||||||
results = await cursor.to_list(length=None)
|
results = await cursor.to_list(length=None)
|
||||||
|
|
||||||
async def update_count(id):
|
async def update_count(id):
|
||||||
query = {"project": {"$eq": id}}
|
query = {"project": {"$eq": id}}
|
||||||
total_count = await db.asset.count_documents(query)
|
total_count = await db.asset.count_documents(query)
|
||||||
update_document = {
|
update_document = {
|
||||||
"$set": {
|
"$set": {
|
||||||
"AssetCount": total_count
|
"AssetCount": total_count
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
await db.project.update_one({"_id": ObjectId(id)}, update_document)
|
||||||
await db.project.update_one({"_id": ObjectId(id)}, update_document)
|
|
||||||
|
|
||||||
fetch_tasks = [update_count(r['id']) for r in results]
|
fetch_tasks = [update_count(r['id']) for r in results]
|
||||||
|
|
||||||
await asyncio.gather(*fetch_tasks)
|
await asyncio.gather(*fetch_tasks)
|
||||||
|
|
||||||
|
|
||||||
@router.post("/project/content")
|
@router.post("/project/content")
|
||||||
|
|
|
@ -256,7 +256,8 @@ async def get_sensitive_result_data2(request_data: dict, db=Depends(get_mongo_db
|
||||||
"time": 1,
|
"time": 1,
|
||||||
"sid": 1,
|
"sid": 1,
|
||||||
"match": 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",
|
"_id": "$url",
|
||||||
"time": {"$first": "$time"}, # 记录相同url下最早插入数据的时间
|
"time": {"$first": "$time"}, # 记录相同url下最早插入数据的时间
|
||||||
"url": {"$first": "$url"},
|
"url": {"$first": "$url"},
|
||||||
"body_id": {"$last": {"$toString": "$_id"}}, # 记录相同url下最早插入数据的_id
|
"body_id": {"$last": {"$toString": "$md5"}}, # 记录相同url下最早插入数据的_id
|
||||||
"children": {
|
"children": {
|
||||||
"$push": {
|
"$push": {
|
||||||
"id": {"$toString": "$_id"},
|
"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}
|
return {"message": "ID is missing in the request data", "code": 400}
|
||||||
|
|
||||||
# Query the database for content based on ID
|
# 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)
|
doc = await db.SensitiveResult.find_one(query)
|
||||||
|
|
||||||
if not doc:
|
if not doc:
|
||||||
|
|
|
@ -8,7 +8,7 @@ import string
|
||||||
|
|
||||||
import yaml
|
import yaml
|
||||||
|
|
||||||
VERSION = "1.3"
|
VERSION = "1.4"
|
||||||
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"
|
||||||
|
|
|
@ -94,7 +94,7 @@ async def create_database():
|
||||||
# {"name": "DomainDic", 'value': domainDict, 'type': 'domainDict'})
|
# {"name": "DomainDic", 'value': domainDict, 'type': 'domainDict'})
|
||||||
sensitive_data = get_sensitive()
|
sensitive_data = get_sensitive()
|
||||||
collection = client[DATABASE_NAME]["SensitiveRule"]
|
collection = client[DATABASE_NAME]["SensitiveRule"]
|
||||||
if sensitiveList:
|
if sensitive_data:
|
||||||
await collection.insert_many(sensitive_data)
|
await collection.insert_many(sensitive_data)
|
||||||
|
|
||||||
collection = client[DATABASE_NAME]["ScheduledTasks"]
|
collection = client[DATABASE_NAME]["ScheduledTasks"]
|
||||||
|
|
File diff suppressed because it is too large
Load Diff
95
main.py
95
main.py
|
@ -8,7 +8,7 @@ from starlette.middleware.base import BaseHTTPMiddleware
|
||||||
from starlette.staticfiles import StaticFiles
|
from starlette.staticfiles import StaticFiles
|
||||||
|
|
||||||
from core.config import *
|
from core.config import *
|
||||||
from core.default import get_dirDict, get_domainDict
|
from core.default import get_dirDict, get_domainDict, get_sensitive
|
||||||
|
|
||||||
set_config()
|
set_config()
|
||||||
|
|
||||||
|
@ -35,43 +35,64 @@ from core.apscheduler_handler import scheduler
|
||||||
|
|
||||||
async def update():
|
async def update():
|
||||||
async for db in get_mongo_db():
|
async for db in get_mongo_db():
|
||||||
# 默认项目有个root_domain为空导致匹配上所有资产
|
# 判断版本
|
||||||
cursor = db.project.find({"root_domains": ""}, {"_id": 1, "root_domains": 1})
|
result = await db.config.find_one({"name": "version"})
|
||||||
async for document in cursor:
|
version = 0
|
||||||
logger.info("Update found empty root_domains")
|
update = False
|
||||||
root_domain = []
|
if result is not None:
|
||||||
for root in document["root_domains"]:
|
version = result["version"]
|
||||||
if root != "":
|
update = result["update"]
|
||||||
root_domain.append(root)
|
if version < float(VERSION):
|
||||||
update_document = {
|
update = False
|
||||||
"$set": {
|
else:
|
||||||
"root_domains": root_domain,
|
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)
|
||||||
await db.project.update_one({"_id": document['_id']}, update_document)
|
# 修改目录字典存储方式
|
||||||
# 修改目录字典存储方式
|
fs = AsyncIOMotorGridFSBucket(db)
|
||||||
fs = AsyncIOMotorGridFSBucket(db)
|
result = await db.config.find_one({"name": "DirDic"})
|
||||||
result = await db.config.find_one({"name": "DirDic"})
|
if result:
|
||||||
if result:
|
await db.config.delete_one({"name": "DirDic"})
|
||||||
await db.config.delete_one({"name": "DirDic"})
|
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('dirdict', byte_content)
|
logger.info("Document DirDict uploaded to GridFS.")
|
||||||
logger.info("Document DirDict uploaded to GridFS.")
|
else:
|
||||||
else:
|
logger.error("No dirdict content to upload.")
|
||||||
logger.error("No dirdict content to upload.")
|
# 修改子域名字典存储方式
|
||||||
# 修改子域名字典存储方式
|
result = await db.config.find_one({"name": "DomainDic"})
|
||||||
result = await db.config.find_one({"name": "DomainDic"})
|
if result:
|
||||||
if result:
|
await db.config.delete_one({"name": "DomainDic"})
|
||||||
await db.config.delete_one({"name": "DomainDic"})
|
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('DomainDic', byte_content)
|
logger.info("Document DomainDic uploaded to GridFS.")
|
||||||
logger.info("Document DomainDic uploaded to GridFS.")
|
else:
|
||||||
else:
|
logger.error("No DomainDic content to upload.")
|
||||||
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")
|
@app.on_event("startup")
|
||||||
|
|
Loading…
Reference in New Issue