From 14cb43711242b5ec531ee68b68fcc371cd23dab0 Mon Sep 17 00:00:00 2001
From: Marvin Zhang <tikazyq@gmail.com>
Date: Wed, 12 Jun 2019 21:21:32 +0800
Subject: [PATCH] add auth to mongodb connection

---
 crawlab/config/config.py      | 11 +++++++----
 crawlab/db/manager.py         |  8 ++++++--
 crawlab/spiders/spiders/db.py | 12 ++++++++----
 crawlab/tasks/spider.py       |  5 ++++-
 4 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/crawlab/config/config.py b/crawlab/config/config.py
index da185c57..8702e00f 100644
--- a/crawlab/config/config.py
+++ b/crawlab/config/config.py
@@ -14,6 +14,13 @@ PROJECT_LOGS_FOLDER = '/var/log/crawlab'
 # 打包临时文件夹
 PROJECT_TMP_FOLDER = '/tmp'
 
+# MongoDB 变量
+MONGO_HOST = '127.0.0.1'
+MONGO_PORT = 27017
+MONGO_USERNAME = None
+MONGO_PASSWORD = None
+MONGO_DB = 'crawlab_test'
+
 # Celery中间者URL
 BROKER_URL = 'redis://127.0.0.1:6379/0'
 
@@ -35,10 +42,6 @@ CELERY_ENABLE_UTC = True
 # flower variables
 FLOWER_API_ENDPOINT = 'http://localhost:5555/api'
 
-# MongoDB 变量
-MONGO_HOST = '127.0.0.1'
-MONGO_PORT = 27017
-MONGO_DB = 'crawlab_test'
 
 # Flask 变量
 DEBUG = False
diff --git a/crawlab/db/manager.py b/crawlab/db/manager.py
index e80aada6..17d6b1ae 100644
--- a/crawlab/db/manager.py
+++ b/crawlab/db/manager.py
@@ -1,7 +1,7 @@
 from bson import ObjectId
 from mongoengine import connect
 from pymongo import MongoClient, DESCENDING
-from config import MONGO_HOST, MONGO_PORT, MONGO_DB
+from config import MONGO_HOST, MONGO_PORT, MONGO_DB, MONGO_USERNAME, MONGO_PASSWORD
 from utils import is_object_id
 
 connect(db=MONGO_DB, host=MONGO_HOST, port=MONGO_PORT)
@@ -13,7 +13,11 @@ class DbManager(object):
     """
 
     def __init__(self):
-        self.mongo = MongoClient(host=MONGO_HOST, port=MONGO_PORT, connect=False)
+        self.mongo = MongoClient(host=MONGO_HOST,
+                                 port=MONGO_PORT,
+                                 username=MONGO_USERNAME,
+                                 password=MONGO_PASSWORD,
+                                 connect=False)
         self.db = self.mongo[MONGO_DB]
 
     def save(self, col_name: str, item: dict, **kwargs) -> None:
diff --git a/crawlab/spiders/spiders/db.py b/crawlab/spiders/spiders/db.py
index 18925f8d..0521d784 100644
--- a/crawlab/spiders/spiders/db.py
+++ b/crawlab/spiders/spiders/db.py
@@ -2,11 +2,15 @@ import os
 
 from pymongo import MongoClient
 
-MONGO_HOST = os.environ.get('MONGO_HOST')
-MONGO_PORT = int(os.environ.get('MONGO_PORT'))
-MONGO_DB = os.environ.get('MONGO_DB')
+MONGO_HOST = os.environ.get('MONGO_HOST') or 'localhost'
+MONGO_PORT = int(os.environ.get('MONGO_PORT')) or 27017
+MONGO_USERNAME = os.environ.get('MONGO_USERNAME')
+MONGO_PASSWORD = os.environ.get('MONGO_PASSWORD')
+MONGO_DB = os.environ.get('MONGO_DB') or 'crawlab_test'
 mongo = MongoClient(host=MONGO_HOST,
-                    port=MONGO_PORT)
+                    port=MONGO_PORT,
+                    username=MONGO_USERNAME,
+                    password=MONGO_PASSWORD)
 db = mongo[MONGO_DB]
 task_id = os.environ.get('CRAWLAB_TASK_ID')
 col_name = os.environ.get('CRAWLAB_COLLECTION')
diff --git a/crawlab/tasks/spider.py b/crawlab/tasks/spider.py
index a568dc0f..e03bab66 100644
--- a/crawlab/tasks/spider.py
+++ b/crawlab/tasks/spider.py
@@ -6,7 +6,8 @@ from time import sleep
 from bson import ObjectId
 from pymongo import ASCENDING, DESCENDING
 
-from config import PROJECT_DEPLOY_FILE_FOLDER, PROJECT_LOGS_FOLDER, MONGO_HOST, MONGO_PORT, MONGO_DB
+from config import PROJECT_DEPLOY_FILE_FOLDER, PROJECT_LOGS_FOLDER, MONGO_HOST, MONGO_PORT, MONGO_DB, MONGO_USERNAME, \
+    MONGO_PASSWORD
 from constants.task import TaskStatus
 from db.manager import db_manager
 from .celery import celery_app
@@ -212,6 +213,8 @@ def execute_config_spider(self, id: str, params: str = None):
     env['MONGO_HOST'] = MONGO_HOST
     env['MONGO_PORT'] = str(MONGO_PORT)
     env['MONGO_DB'] = MONGO_DB
+    env['MONGO_USERNAME'] = MONGO_USERNAME
+    env['MONGO_PASSWORD'] = MONGO_PASSWORD
 
     cmd_arr = [
         sys.executable,