2024-06-05 13:39:34 +00:00
|
|
|
|
# -*- coding:utf-8 -*-
|
|
|
|
|
# @name: db
|
|
|
|
|
# @auth: rainy-autumn@outlook.com
|
|
|
|
|
# @version:
|
2024-06-13 14:11:11 +00:00
|
|
|
|
import time
|
2024-06-11 13:57:44 +00:00
|
|
|
|
from urllib.parse import quote_plus
|
2024-07-20 09:23:44 +00:00
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorGridFSBucket
|
2024-06-05 13:39:34 +00:00
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCursor
|
|
|
|
|
from core.default import *
|
|
|
|
|
from core.config import *
|
|
|
|
|
from loguru import logger
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def get_mongo_db():
|
2024-07-17 13:39:09 +00:00
|
|
|
|
|
|
|
|
|
client = AsyncIOMotorClient(f"mongodb://{DATABASE_USER}:{quote_plus(DATABASE_PASSWORD)}@{MONGODB_IP}:{str(MONGODB_PORT)}",
|
2024-06-05 13:39:34 +00:00
|
|
|
|
serverSelectionTimeoutMS=10000, unicode_decode_error_handler='ignore')
|
|
|
|
|
db = client[DATABASE_NAME]
|
|
|
|
|
try:
|
|
|
|
|
yield db
|
|
|
|
|
finally:
|
|
|
|
|
client.close()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def create_database():
|
|
|
|
|
client = None
|
2024-06-13 14:11:11 +00:00
|
|
|
|
check_flag = 0
|
2024-06-05 13:39:34 +00:00
|
|
|
|
try:
|
2024-06-13 14:11:11 +00:00
|
|
|
|
while True:
|
|
|
|
|
try:
|
|
|
|
|
# 创建新的 MongoDB 客户端
|
|
|
|
|
client = AsyncIOMotorClient(f"mongodb://{quote_plus(DATABASE_USER)}:{quote_plus(DATABASE_PASSWORD)}@{MONGODB_IP}:{str(MONGODB_PORT)}",
|
|
|
|
|
serverSelectionTimeoutMS=2000)
|
|
|
|
|
break
|
|
|
|
|
except Exception as e:
|
2024-06-30 15:40:24 +00:00
|
|
|
|
time.sleep(10)
|
2024-06-13 14:11:11 +00:00
|
|
|
|
check_flag += 1
|
2024-06-30 15:40:24 +00:00
|
|
|
|
if check_flag == 10:
|
2024-06-13 14:11:11 +00:00
|
|
|
|
logger.error(f"Error re creating database: {e}")
|
2024-06-30 15:40:24 +00:00
|
|
|
|
exit(1)
|
2024-06-05 13:39:34 +00:00
|
|
|
|
# 获取数据库列表
|
|
|
|
|
database_names = await client.list_database_names()
|
2024-07-20 09:45:47 +00:00
|
|
|
|
db = client[DATABASE_NAME]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
# 如果数据库不存在,创建数据库
|
|
|
|
|
if DATABASE_NAME not in database_names:
|
|
|
|
|
# 在数据库中创建一个集合,比如名为 "user"
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["user"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
|
|
|
|
|
# 用户数据
|
|
|
|
|
await collection.insert_one({"username": "ScopeSentry",
|
|
|
|
|
'password': 'b0ce71fcbed8a6ca579d52800145119cc7d999dc8651b62dfc1ced9a984e6e64'})
|
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["config"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
# 系统配置
|
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "timezone", 'value': 'Asia/Shanghai', 'type': 'system'})
|
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "MaxTaskNum", 'value': '7', 'type': 'system'})
|
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "DirscanThread", 'value': '15', 'type': 'system'})
|
|
|
|
|
await collection.insert_one(
|
2024-07-14 10:09:30 +00:00
|
|
|
|
{"name": "PortscanThread", 'value': '5', 'type': 'system'})
|
2024-06-05 13:39:34 +00:00
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "CrawlerThread", 'value': '2', 'type': 'system'})
|
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "UrlMaxNum", 'value': '500', 'type': 'system'})
|
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "UrlThread", 'value': '5', 'type': 'system'})
|
|
|
|
|
# 设置时区为Asia/Shanghai
|
|
|
|
|
# SHA_TZ = timezone(TIMEZONE)
|
|
|
|
|
# timezone('Asia/Shanghai')
|
|
|
|
|
# utc_now = datetime.utcnow().replace(tzinfo=timezone.utc)
|
|
|
|
|
# time_now = utc_now.astimezone(SHA_TZ)
|
|
|
|
|
# formatted_time = time_now.strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
|
# subfinder配置
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["config"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
# 插入一条数据
|
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "SubfinderApiConfig", 'value': subfinderApiConfig, 'type': 'subfinder'})
|
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "RadConfig", 'value': radConfig, 'type': 'rad'})
|
2024-07-07 15:21:34 +00:00
|
|
|
|
# dirDict = get_dirDict()
|
|
|
|
|
# await collection.insert_one(
|
|
|
|
|
# {"name": "DirDic", 'value': dirDict, 'type': 'dirDict'})
|
2024-07-20 09:23:44 +00:00
|
|
|
|
# 目录扫描字典
|
2024-07-20 09:45:47 +00:00
|
|
|
|
fs = AsyncIOMotorGridFSBucket(db)
|
2024-07-20 09:23:44 +00:00
|
|
|
|
content = get_dirDict()
|
|
|
|
|
if content:
|
|
|
|
|
byte_content = content.encode('utf-8')
|
|
|
|
|
await fs.upload_from_stream('dirdict', byte_content)
|
|
|
|
|
# 子域名字典
|
|
|
|
|
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.")
|
|
|
|
|
|
2024-06-05 13:39:34 +00:00
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"name": "notification", 'dirScanNotification': True,
|
|
|
|
|
'portScanNotification': True, 'sensitiveNotification': True,
|
|
|
|
|
'subdomainTakeoverNotification': True,
|
|
|
|
|
'pageMonNotification': True,
|
|
|
|
|
'subdomainNotification': True,
|
|
|
|
|
'vulNotification': True,
|
|
|
|
|
'type': 'notification'})
|
2024-07-07 15:21:34 +00:00
|
|
|
|
# domainDict = get_domainDict()
|
|
|
|
|
# await collection.insert_one(
|
|
|
|
|
# {"name": "DomainDic", 'value': domainDict, 'type': 'domainDict'})
|
2024-06-05 13:39:34 +00:00
|
|
|
|
sensitive_data = get_sensitive()
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["SensitiveRule"]
|
2024-07-16 15:30:31 +00:00
|
|
|
|
if sensitive_data:
|
2024-06-05 13:39:34 +00:00
|
|
|
|
await collection.insert_many(sensitive_data)
|
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["ScheduledTasks"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"id": "page_monitoring", "name": "Page Monitoring", 'hour': 24, 'node': [], 'allNode': True, 'type': 'Page Monitoring', 'state': True})
|
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
await db.create_collection("notification")
|
2024-06-05 13:39:34 +00:00
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["PortDict"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
await collection.insert_many(portDic)
|
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["PocList"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
pocData = get_poc()
|
|
|
|
|
await collection.insert_many(pocData)
|
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["project"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
project_data, target_data = get_project_data()
|
|
|
|
|
await collection.insert_many(project_data)
|
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["ProjectTargetData"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
await collection.insert_many(target_data)
|
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["FingerprintRules"]
|
2024-06-30 15:40:24 +00:00
|
|
|
|
fingerprint = get_finger()
|
|
|
|
|
await collection.insert_many(fingerprint)
|
2024-06-05 13:39:34 +00:00
|
|
|
|
else:
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["config"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
result = await collection.find_one({"name": "timezone"})
|
|
|
|
|
set_timezone(result.get('value', 'Asia/Shanghai'))
|
|
|
|
|
|
2024-07-20 09:45:47 +00:00
|
|
|
|
collection = db["ScheduledTasks"]
|
2024-06-05 13:39:34 +00:00
|
|
|
|
result = await collection.find_one({"id": "page_monitoring"})
|
|
|
|
|
if not result:
|
|
|
|
|
await collection.insert_one(
|
|
|
|
|
{"id": "page_monitoring", "name": "Page Monitoring", 'hour': 24, 'type': 'Page Monitoring', 'state': True})
|
2024-07-20 09:45:47 +00:00
|
|
|
|
await get_fingerprint(db)
|
|
|
|
|
# await get_sens_rule(db)
|
|
|
|
|
await get_project(db)
|
2024-06-05 13:39:34 +00:00
|
|
|
|
except Exception as e:
|
|
|
|
|
# 处理异常
|
|
|
|
|
logger.error(f"Error creating database: {e}")
|
|
|
|
|
exit(0)
|
|
|
|
|
finally:
|
|
|
|
|
# 在适当的地方关闭 MongoDB 客户端
|
|
|
|
|
if client:
|
|
|
|
|
client.close()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def get_fingerprint(client):
|
|
|
|
|
collection = client["FingerprintRules"]
|
|
|
|
|
cursor = collection.find({}, {"_id": 1, "name": 1})
|
|
|
|
|
async for document in cursor:
|
|
|
|
|
document['id'] = str(document['_id'])
|
|
|
|
|
del document['_id']
|
|
|
|
|
APP[document['id']] = document['name']
|
|
|
|
|
|
|
|
|
|
|
2024-07-14 10:09:30 +00:00
|
|
|
|
# async def get_sens_rule(client):
|
|
|
|
|
# collection = client["SensitiveRule"]
|
|
|
|
|
# cursor = collection.find({}, {"_id": 1, "name": 1, "color": 1})
|
|
|
|
|
# async for document in cursor:
|
|
|
|
|
# document['id'] = str(document['_id'])
|
|
|
|
|
# del document['_id']
|
|
|
|
|
# SensitiveRuleList[document['id']] = {
|
|
|
|
|
# "name": document['name'],
|
|
|
|
|
# "color": document['color']
|
|
|
|
|
# }
|
2024-06-05 13:39:34 +00:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def get_project(client):
|
|
|
|
|
collection = client["project"]
|
|
|
|
|
cursor = collection.find({}, {"_id": 1, "name": 1})
|
|
|
|
|
async for document in cursor:
|
|
|
|
|
document['id'] = str(document['_id'])
|
|
|
|
|
Project_List[document['name'].lower()] = document['id']
|