73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
# @Author : jeffzhang
|
|
# @Time : 18-5-10
|
|
# @File : mongo_db.py
|
|
# @Desc : ""
|
|
|
|
from flask import Flask
|
|
from pymongo import MongoClient
|
|
from instance import config
|
|
|
|
ProductionConfig = config.ProductionConfig
|
|
app = Flask(__name__)
|
|
app.config.from_object(ProductionConfig)
|
|
db_host = app.config.get('DB_HOST')
|
|
db_port = app.config.get('DB_PORT')
|
|
db_username = app.config.get('DB_USERNAME')
|
|
db_password = app.config.get('DB_PASSWORD')
|
|
db_name = app.config.get('DB_NAME')
|
|
|
|
|
|
def connectiondb(collection):
|
|
client = MongoClient(db_host, db_port)
|
|
db = client[db_name]
|
|
db.authenticate(db_username, db_password)
|
|
dbcollection = db[collection]
|
|
return dbcollection
|
|
|
|
|
|
def db_management(command):
|
|
client = MongoClient(db_host, db_port)
|
|
db = client[db_name]
|
|
db.authenticate(db_username, db_password)
|
|
if command == 'collection_names':
|
|
result = db.collection_names()
|
|
return result
|
|
|
|
|
|
def db_name_conf():
|
|
asset_db = app.config.get('ASSET_DB')
|
|
tasks_db = app.config.get('TASKS_DB')
|
|
vul_db = app.config.get('VULNERABILITY_DB')
|
|
plugin_db = app.config.get('PLUGIN_DB')
|
|
config_db = app.config.get('CONFIG_DB')
|
|
server_db = app.config.get('SERVER_DB')
|
|
subdomain_db = app.config.get('SUBDOMAIN_DB')
|
|
domain_db = app.config.get('DOMAIN_DB')
|
|
weekpasswd_db = app.config.get('WEEKPASSWD_DB')
|
|
port_db = app.config.get('PORT_DB')
|
|
auth_db = app.config.get('AUTH_DB')
|
|
search_db = app.config.get('SEARCH_DB')
|
|
acunetix_db = app.config.get('ACUNETIX_DB')
|
|
db_name_dict = {
|
|
'asset_db': asset_db,
|
|
'tasks_db': tasks_db,
|
|
'vul_db': vul_db,
|
|
'plugin_db': plugin_db,
|
|
'config_db': config_db,
|
|
'server_db': server_db,
|
|
'subdomain_db': subdomain_db,
|
|
'domain_db': domain_db,
|
|
'weekpasswd_db': weekpasswd_db,
|
|
'port_db': port_db,
|
|
'auth_db': auth_db,
|
|
'search_db': search_db,
|
|
'acunetix_db': acunetix_db,
|
|
}
|
|
return db_name_dict
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print db_management('collection_names')
|