chore: improve mask/unmask encrypted_extra (#29943)

This commit is contained in:
Beto Dealmeida
2024-08-22 16:45:32 -04:00
committed by GitHub
parent bf94370d38
commit 4b59e42d3f
13 changed files with 490 additions and 151 deletions

View File

@@ -19,6 +19,7 @@
from __future__ import annotations
import json
from textwrap import dedent
from typing import Any
@@ -334,3 +335,60 @@ def test_quote_table() -> None:
BaseEngineSpec.quote_table(Table("ta ble", "sche.ma", 'cata"log'), dialect)
== '"cata""log"."sche.ma"."ta ble"'
)
def test_mask_encrypted_extra() -> None:
"""
Test that the private key is masked when the database is edited.
"""
from superset.db_engine_specs.base import BaseEngineSpec
config = json.dumps(
{
"foo": "bar",
"service_account_info": {
"project_id": "black-sanctum-314419",
"private_key": "SECRET",
},
}
)
assert BaseEngineSpec.mask_encrypted_extra(config) == json.dumps(
{
"foo": "XXXXXXXXXX",
"service_account_info": "XXXXXXXXXX",
}
)
def test_unmask_encrypted_extra() -> None:
"""
Test that the private key can be reused from the previous `encrypted_extra`.
"""
from superset.db_engine_specs.base import BaseEngineSpec
old = json.dumps(
{
"foo": "bar",
"service_account_info": {
"project_id": "black-sanctum-314419",
"private_key": "SECRET",
},
}
)
new = json.dumps(
{
"foo": "XXXXXXXXXX",
"service_account_info": "XXXXXXXXXX",
}
)
assert BaseEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"foo": "bar",
"service_account_info": {
"project_id": "black-sanctum-314419",
"private_key": "SECRET",
},
}
)

View File

@@ -191,7 +191,7 @@ def test_get_parameters_from_uri_serializable() -> None:
def test_unmask_encrypted_extra() -> None:
"""
Test that the private key can be reused from the previous ``encrypted_extra``.
Test that the private key can be reused from the previous `encrypted_extra`.
"""
from superset.db_engine_specs.bigquery import BigQueryEngineSpec
@@ -212,17 +212,52 @@ def test_unmask_encrypted_extra() -> None:
}
)
assert json.loads(str(BigQueryEngineSpec.unmask_encrypted_extra(old, new))) == {
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "SECRET",
},
}
assert BigQueryEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "SECRET",
},
}
)
def test_unmask_encrypted_extra_when_empty() -> None:
def test_unmask_encrypted_extra_field_changeed() -> None:
"""
Test that a None value works for ``encrypted_extra``.
Test that the private key is not reused when the field has changed.
"""
from superset.db_engine_specs.bigquery import BigQueryEngineSpec
old = json.dumps(
{
"credentials_info": {
"project_id": "black-sanctum-314419",
"private_key": "SECRET",
},
}
)
new = json.dumps(
{
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "NEW-SECRET",
},
}
)
assert BigQueryEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "NEW-SECRET",
},
}
)
def test_unmask_encrypted_extra_when_old_is_none() -> None:
"""
Test that a `None` value for the old field works for `encrypted_extra`.
"""
from superset.db_engine_specs.bigquery import BigQueryEngineSpec
@@ -236,17 +271,19 @@ def test_unmask_encrypted_extra_when_empty() -> None:
}
)
assert json.loads(str(BigQueryEngineSpec.unmask_encrypted_extra(old, new))) == {
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "XXXXXXXXXX",
},
}
assert BigQueryEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"credentials_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "XXXXXXXXXX",
},
}
)
def test_unmask_encrypted_extra_when_new_is_empty() -> None:
def test_unmask_encrypted_extra_when_new_is_none() -> None:
"""
Test that a None value works for ``encrypted_extra``.
Test that a `None` value for the new field works for `encrypted_extra`.
"""
from superset.db_engine_specs.bigquery import BigQueryEngineSpec
@@ -263,6 +300,31 @@ def test_unmask_encrypted_extra_when_new_is_empty() -> None:
assert BigQueryEngineSpec.unmask_encrypted_extra(old, new) is None
def test_mask_encrypted_extra() -> None:
"""
Test that the private key is masked when the database is edited.
"""
from superset.db_engine_specs.bigquery import BigQueryEngineSpec
config = json.dumps(
{
"credentials_info": {
"project_id": "black-sanctum-314419",
"private_key": "SECRET",
},
}
)
assert BigQueryEngineSpec.mask_encrypted_extra(config) == json.dumps(
{
"credentials_info": {
"project_id": "black-sanctum-314419",
"private_key": "XXXXXXXXXX",
},
}
)
def test_mask_encrypted_extra_when_empty() -> None:
"""
Test that the encrypted extra will return a none value if the field is empty.

View File

@@ -247,9 +247,34 @@ def test_validate_parameters_catalog_and_credentials(
)
def test_mask_encrypted_extra() -> None:
"""
Test that the private key is masked when the database is edited.
"""
from superset.db_engine_specs.gsheets import GSheetsEngineSpec
config = json.dumps(
{
"service_account_info": {
"project_id": "black-sanctum-314419",
"private_key": "SECRET",
},
}
)
assert GSheetsEngineSpec.mask_encrypted_extra(config) == json.dumps(
{
"service_account_info": {
"project_id": "black-sanctum-314419",
"private_key": "XXXXXXXXXX",
},
}
)
def test_unmask_encrypted_extra() -> None:
"""
Test that the private key can be reused from the previous ``encrypted_extra``.
Test that the private key can be reused from the previous `encrypted_extra`.
"""
from superset.db_engine_specs.gsheets import GSheetsEngineSpec
@@ -270,17 +295,52 @@ def test_unmask_encrypted_extra() -> None:
}
)
assert json.loads(str(GSheetsEngineSpec.unmask_encrypted_extra(old, new))) == {
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "SECRET",
},
}
assert GSheetsEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "SECRET",
},
}
)
def test_unmask_encrypted_extra_field_changeed() -> None:
"""
Test that the private key is not reused when the field has changed.
"""
from superset.db_engine_specs.gsheets import GSheetsEngineSpec
old = json.dumps(
{
"service_account_info": {
"project_id": "black-sanctum-314419",
"private_key": "SECRET",
},
}
)
new = json.dumps(
{
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "NEW-SECRET",
},
}
)
assert GSheetsEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "NEW-SECRET",
},
}
)
def test_unmask_encrypted_extra_when_old_is_none() -> None:
"""
Test that a None value works for ``encrypted_extra``.
Test that a `None` value for the old field works for `encrypted_extra`.
"""
from superset.db_engine_specs.gsheets import GSheetsEngineSpec
@@ -294,17 +354,19 @@ def test_unmask_encrypted_extra_when_old_is_none() -> None:
}
)
assert json.loads(str(GSheetsEngineSpec.unmask_encrypted_extra(old, new))) == {
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "XXXXXXXXXX",
},
}
assert GSheetsEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"service_account_info": {
"project_id": "yellow-unicorn-314419",
"private_key": "XXXXXXXXXX",
},
}
)
def test_unmask_encrypted_extra_when_new_is_none() -> None:
"""
Test that a None value works for ``encrypted_extra``.
Test that a `None` value for the new field works for `encrypted_extra`.
"""
from superset.db_engine_specs.gsheets import GSheetsEngineSpec

View File

@@ -291,3 +291,106 @@ def test_get_default_catalog() -> None:
sqlalchemy_uri="snowflake://user:pass@account/database_name/default",
)
assert SnowflakeEngineSpec.get_default_catalog(database) == "database_name"
def test_mask_encrypted_extra() -> None:
"""
Test that the private keys are masked when the database is edited.
"""
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
config = json.dumps(
{
"auth_method": "keypair",
"auth_params": {
"privatekey_body": (
"-----BEGIN ENCRYPTED PRIVATE KEY-----"
"..."
"-----END ENCRYPTED PRIVATE KEY-----"
),
"privatekey_pass": "my_password",
},
}
)
assert SnowflakeEngineSpec.mask_encrypted_extra(config) == json.dumps(
{
"auth_method": "keypair",
"auth_params": {
"privatekey_body": "XXXXXXXXXX",
"privatekey_pass": "XXXXXXXXXX",
},
}
)
def test_mask_encrypted_extra_no_fields() -> None:
"""
Test that the private key is masked when the database is edited.
"""
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
config = json.dumps(
{
# this is a fake example and the fields are made up
"auth_method": "token",
"auth_params": {
"jwt": "SECRET",
},
}
)
assert SnowflakeEngineSpec.mask_encrypted_extra(config) == json.dumps(
{
"auth_method": "token",
"auth_params": {
"jwt": "SECRET",
},
}
)
def test_unmask_encrypted_extra() -> None:
"""
Test that the private keys can be reused from the previous `encrypted_extra`.
"""
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
old = json.dumps(
{
"auth_method": "keypair",
"auth_params": {
"privatekey_body": (
"-----BEGIN ENCRYPTED PRIVATE KEY-----"
"..."
"-----END ENCRYPTED PRIVATE KEY-----"
),
"privatekey_pass": "my_password",
},
}
)
new = json.dumps(
{
"foo": "bar",
"auth_method": "keypair",
"auth_params": {
"privatekey_body": "XXXXXXXXXX",
"privatekey_pass": "XXXXXXXXXX",
},
}
)
assert SnowflakeEngineSpec.unmask_encrypted_extra(old, new) == json.dumps(
{
"foo": "bar",
"auth_method": "keypair",
"auth_params": {
"privatekey_body": (
"-----BEGIN ENCRYPTED PRIVATE KEY-----"
"..."
"-----END ENCRYPTED PRIVATE KEY-----"
),
"privatekey_pass": "my_password",
},
}
)