update init db

This commit is contained in:
Autumn.home 2024-07-20 17:45:47 +08:00
parent 2c6069a913
commit 6164f5aaee
1 changed files with 18 additions and 19 deletions

View File

@ -40,17 +40,17 @@ async def create_database():
exit(1) exit(1)
# 获取数据库列表 # 获取数据库列表
database_names = await client.list_database_names() database_names = await client.list_database_names()
db = client[DATABASE_NAME]
# 如果数据库不存在,创建数据库 # 如果数据库不存在,创建数据库
if DATABASE_NAME not in database_names: if DATABASE_NAME not in database_names:
# 在数据库中创建一个集合,比如名为 "user" # 在数据库中创建一个集合,比如名为 "user"
collection = client[DATABASE_NAME]["user"] collection = db["user"]
# 用户数据 # 用户数据
await collection.insert_one({"username": "ScopeSentry", await collection.insert_one({"username": "ScopeSentry",
'password': 'b0ce71fcbed8a6ca579d52800145119cc7d999dc8651b62dfc1ced9a984e6e64'}) 'password': 'b0ce71fcbed8a6ca579d52800145119cc7d999dc8651b62dfc1ced9a984e6e64'})
collection = client[DATABASE_NAME]["config"] collection = db["config"]
# 系统配置 # 系统配置
await collection.insert_one( await collection.insert_one(
{"name": "timezone", 'value': 'Asia/Shanghai', 'type': 'system'}) {"name": "timezone", 'value': 'Asia/Shanghai', 'type': 'system'})
@ -73,7 +73,7 @@ async def create_database():
# time_now = utc_now.astimezone(SHA_TZ) # time_now = utc_now.astimezone(SHA_TZ)
# formatted_time = time_now.strftime("%Y-%m-%d %H:%M:%S") # formatted_time = time_now.strftime("%Y-%m-%d %H:%M:%S")
# subfinder配置 # subfinder配置
collection = client[DATABASE_NAME]["config"] collection = db["config"]
# 插入一条数据 # 插入一条数据
await collection.insert_one( await collection.insert_one(
{"name": "SubfinderApiConfig", 'value': subfinderApiConfig, 'type': 'subfinder'}) {"name": "SubfinderApiConfig", 'value': subfinderApiConfig, 'type': 'subfinder'})
@ -83,7 +83,7 @@ async def create_database():
# await collection.insert_one( # await collection.insert_one(
# {"name": "DirDic", 'value': dirDict, 'type': 'dirDict'}) # {"name": "DirDic", 'value': dirDict, 'type': 'dirDict'})
# 目录扫描字典 # 目录扫描字典
fs = AsyncIOMotorGridFSBucket(client) fs = AsyncIOMotorGridFSBucket(db)
content = get_dirDict() content = get_dirDict()
if content: if content:
byte_content = content.encode('utf-8') byte_content = content.encode('utf-8')
@ -107,47 +107,46 @@ async def create_database():
# await collection.insert_one( # await collection.insert_one(
# {"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 = db["SensitiveRule"]
if sensitive_data: if sensitive_data:
await collection.insert_many(sensitive_data) await collection.insert_many(sensitive_data)
collection = client[DATABASE_NAME]["ScheduledTasks"] collection = db["ScheduledTasks"]
await collection.insert_one( await collection.insert_one(
{"id": "page_monitoring", "name": "Page Monitoring", 'hour': 24, 'node': [], 'allNode': True, 'type': 'Page Monitoring', 'state': True}) {"id": "page_monitoring", "name": "Page Monitoring", 'hour': 24, 'node': [], 'allNode': True, 'type': 'Page Monitoring', 'state': True})
collection = client[DATABASE_NAME] await db.create_collection("notification")
await collection.create_collection("notification")
collection = client[DATABASE_NAME]["PortDict"] collection = db["PortDict"]
await collection.insert_many(portDic) await collection.insert_many(portDic)
collection = client[DATABASE_NAME]["PocList"] collection = db["PocList"]
pocData = get_poc() pocData = get_poc()
await collection.insert_many(pocData) await collection.insert_many(pocData)
collection = client[DATABASE_NAME]["project"] collection = db["project"]
project_data, target_data = get_project_data() project_data, target_data = get_project_data()
await collection.insert_many(project_data) await collection.insert_many(project_data)
collection = client[DATABASE_NAME]["ProjectTargetData"] collection = db["ProjectTargetData"]
await collection.insert_many(target_data) await collection.insert_many(target_data)
collection = client[DATABASE_NAME]["FingerprintRules"] collection = db["FingerprintRules"]
fingerprint = get_finger() fingerprint = get_finger()
await collection.insert_many(fingerprint) await collection.insert_many(fingerprint)
else: else:
collection = client[DATABASE_NAME]["config"] collection = db["config"]
result = await collection.find_one({"name": "timezone"}) result = await collection.find_one({"name": "timezone"})
set_timezone(result.get('value', 'Asia/Shanghai')) set_timezone(result.get('value', 'Asia/Shanghai'))
collection = client[DATABASE_NAME]["ScheduledTasks"] collection = db["ScheduledTasks"]
result = await collection.find_one({"id": "page_monitoring"}) result = await collection.find_one({"id": "page_monitoring"})
if not result: if not result:
await collection.insert_one( await collection.insert_one(
{"id": "page_monitoring", "name": "Page Monitoring", 'hour': 24, 'type': 'Page Monitoring', 'state': True}) {"id": "page_monitoring", "name": "Page Monitoring", 'hour': 24, 'type': 'Page Monitoring', 'state': True})
await get_fingerprint(client[DATABASE_NAME]) await get_fingerprint(db)
# await get_sens_rule(client[DATABASE_NAME]) # await get_sens_rule(db)
await get_project(client[DATABASE_NAME]) await get_project(db)
except Exception as e: except Exception as e:
# 处理异常 # 处理异常
logger.error(f"Error creating database: {e}") logger.error(f"Error creating database: {e}")