Configurable SSL Connection (Netflix#373)

* [TRIS-297] Configurable SSL Connection (#1)

* Configurable SSL connection

* Update services/utils/__init__.py

* no ssl unit testing (#3)

* ssl seperate test (Netflix#4)

* dsn generator sslmode none (Netflix#5)

Author: RikishK

run_goose.py

@@ -51,13 +51,27 @@ def main():
     parser.add_argument("--wait", type=int, default=30, help="Wait for connection for X seconds")
     args = parser.parse_args()
 
-    db_connection_string = "postgresql://{}:{}@{}:{}/{}?sslmode=disable".format(
-        quote(os.environ["MF_METADATA_DB_USER"]),
-        quote(os.environ["MF_METADATA_DB_PSWD"]),
-        os.environ["MF_METADATA_DB_HOST"],
-        os.environ["MF_METADATA_DB_PORT"],
-        os.environ["MF_METADATA_DB_NAME"],
-    )
+    db_connection_string = f'postgresql://{quote(os.environ["MF_METADATA_DB_USER"])}:'\
+        f'{quote(os.environ["MF_METADATA_DB_PSWD"])}@{os.environ["MF_METADATA_DB_HOST"]}:'\
+        f'{os.environ["MF_METADATA_DB_PORT"]}/{os.environ["MF_METADATA_DB_NAME"]}'
+
+    ssl_mode = os.environ["MF_METADATA_DB_SSL_MODE"]
+    ssl_cert_path = os.environ["MF_METADATA_DB_SSL_CERT_PATH"]
+    ssl_key_path = os.environ["MF_METADATA_DB_SSL_KEY_PATH"]
+    ssl_root_cert_path = os.environ["MF_METADATA_DB_SSL_ROOT_CERT"]
+
+    if (ssl_mode in ['allow', 'prefer', 'require', 'verify-ca', 'verify-full']):
+        ssl_query = f'ssl_mode={ssl_mode}'
+        if ssl_cert_path is not None:
+            ssl_query = f'{ssl_query}&sslcert={ssl_cert_path}'
+        if ssl_key_path is not None:
+            ssl_query = f'{ssl_query}&sslkey={ssl_key_path}'
+        if ssl_root_cert_path is not None:
+            ssl_query = f'{ssl_query}&sslrootcert={ssl_root_cert_path}'
+    else:
+        ssl_query = f'sslmode=disable'
+
+    db_connection_string = f'{db_connection_string}?{ssl_query}'
 
     if args.wait:
         wait_for_postgres(db_connection_string, timeout_seconds=args.wait)

services/utils/__init__.py

@@ -193,6 +193,10 @@ def __init__(self,
                  user: str = "postgres",
                  password: str = "postgres",
                  database_name: str = "postgres",
+                 ssl_mode: str = "disabled",
+                 ssl_cert_path: str = None,
+                 ssl_key_path: str = None,
+                 ssl_root_cert_path: str = None,
                  prefix="MF_METADATA_DB_",
                  pool_min: int = 1,
                  pool_max: int = 10,
@@ -209,6 +213,10 @@ def __init__(self,
         self._user = os.environ.get(prefix + "USER", user)
         self._password = os.environ.get(prefix + "PSWD", password)
         self._database_name = os.environ.get(prefix + "NAME", database_name)
+        self._ssl_mode = os.environ.get(prefix + "SSL_MODE", ssl_mode)
+        self._ssl_cert_path = os.environ.get(prefix + "SSL_CERT_PATH", ssl_cert_path)
+        self._ssl_key_path = os.environ.get(prefix + "SSL_KEY_PATH", ssl_key_path),
+        self._ssl_root_cert_path = os.environ.get(prefix + "SSL_ROOT_CERT_PATH", ssl_root_cert_path)
         conn_str_required_values = [
             self._host,
             self._port,
@@ -247,19 +255,44 @@ def _is_valid_dsn(dsn):
 
     @property
     def connection_string_url(self):
-        # postgresql://[user[:password]@][host][:port][/dbname][?param1=value1&...]
-        return f'postgresql://{quote(self._user)}:{quote(self._password)}@{self._host}:{self._port}/{self._database_name}?sslmode=disable'
+        base_url = f'postgresql://{quote(self._user)}:{quote(self._password)}@{self._host}:{self._port}/{self._database_name}'
+        if (self._ssl_mode in ['allow', 'prefer', 'require', 'verify-ca', 'verify-full']):
+            ssl_query = f'sslmode={self._ssl_mode}'
+            if self._ssl_cert_path is not None:
+                ssl_query = f'{ssl_query}&sslcert={self._ssl_cert_path}'
+            if self._ssl_key_path is not None:
+                ssl_query = f'{ssl_query}&sslkey={self._ssl_key_path}'
+            if self._ssl_root_cert_path is not None:
+                ssl_query = f'{ssl_query}&sslrootcert={self._ssl_root_cert_path}'
+        else:
+            ssl_query = f'sslmode=disable'
+
+        return f'{base_url}?{ssl_query}'
 
     @property
     def dsn(self):
         if self._dsn is None:
-            return psycopg2.extensions.make_dsn(
-                dbname=self._database_name,
-                user=self._user,
-                host=self._host,
-                port=self._port,
-                password=self._password
-            )
+            ssl_mode = self._ssl_mode
+            sslcert = self._ssl_cert_path
+            sslkey = self._ssl_key_path
+            sslrootcert = self._ssl_root_cert_path
+            if (ssl_mode not in ['allow', 'prefer', 'require', 'verify-ca', 'verify-full']):
+                ssl_mode = None
+                sslcert = None
+                sslkey = None
+                sslrootcert = None
+            kwargs = {
+                'dbname': self._database_name,
+                'user': self._user,
+                'host': self._host,
+                'port': self._port,
+                'password': self._password,
+                'sslmode': ssl_mode,
+                'sslcert': sslcert,
+                'sslkey': sslkey,
+                'sslrootcert': sslrootcert
+            }
+            return psycopg2.extensions.make_dsn(**{k: v for k, v in kwargs.items() if v is not None})
         else:
             return self._dsn
 

services/utils/tests/unit_tests/utils_test.py

@@ -125,7 +125,10 @@ def test_db_conf_env_dsn():
     with set_env({'MF_METADATA_DB_DSN': 'dbname=testgres user=test_user host=test_host port=1234 password=test_pwd'}):
         # valid DSN in env should set correctly.
         assert DBConfiguration().dsn == 'dbname=testgres user=test_user host=test_host port=1234 password=test_pwd'
-
+        
+    with set_env({'MF_METADATA_DB_DSN': 'dbname=testgres user=test_user host=test_host port=1234 password=test_pwd sslmode=verify-full sslrootcert=/test'}):
+        # valid DSN in env should set correctly.
+        assert DBConfiguration().dsn == 'dbname=testgres user=test_user host=test_host port=1234 password=test_pwd sslmode=verify-full sslrootcert=/test'
 
 def test_db_conf_pool_size():
     with set_env():

tox.ini

@@ -17,3 +17,4 @@ commands = pytest --cov=services -m unit_tests
 
 [testenv:integration]
 commands = pytest --cov=services -m integration_tests
+