local-run: add option to set service auth token in container environment

paasta_tools/cli/cmds/local_run.py

@@ -38,6 +38,7 @@
 from paasta_tools.cli.cmds.cook_image import paasta_cook_image
 from paasta_tools.cli.utils import figure_out_service_name
 from paasta_tools.cli.utils import get_instance_config
+from paasta_tools.cli.utils import get_service_auth_token
 from paasta_tools.cli.utils import lazy_choices_completer
 from paasta_tools.cli.utils import list_instances
 from paasta_tools.cli.utils import pick_random_port
@@ -506,6 +507,17 @@ def add_subparser(subparsers):
         required=False,
         default=False,
     )
+    list_parser.add_argument(
+        "--use-service-auth-token",
+        help=(
+            "Acquire service authentication token for the underlying instance,"
+            " and set it in the container environment"
+        ),
+        action="store_true",
+        dest="use_service_auth_token",
+        required=False,
+        default=False,
+    )
     list_parser.add_argument(
         "--sha",
         help=(
@@ -817,6 +829,7 @@ def run_docker_container(
     assume_role_arn="",
     use_okta_role=False,
     assume_role_aws_account: Optional[str] = None,
+    use_service_auth_token: bool = False,
 ):
     """docker-py has issues running a container with a TTY attached, so for
     consistency we execute 'docker run' directly in both interactive and
@@ -906,6 +919,9 @@ def run_docker_container(
         )
         environment.update(aws_creds)
 
+    if use_service_auth_token:
+        environment["YELP_SVC_AUTHZ_TOKEN"] = get_service_auth_token()
+
     local_run_environment = get_local_run_environment_vars(
         instance_config=instance_config, port0=chosen_port, framework=framework
     )
@@ -1251,6 +1267,7 @@ def configure_and_run_docker_container(
         assume_role_arn=args.assume_role_arn,
         assume_role_aws_account=assume_role_aws_account,
         use_okta_role=args.use_okta_role,
+        use_service_auth_token=args.use_service_auth_token,
     )
 
 

paasta_tools/cli/utils.py

@@ -37,6 +37,8 @@
 from typing import Tuple
 
 import ephemeral_port_reserve
+from botocore.credentials import InstanceMetadataFetcher
+from botocore.credentials import InstanceMetadataProvider
 from mypy_extensions import NamedArg
 
 from paasta_tools import remote_git
@@ -76,6 +78,22 @@
 from paasta_tools.utils import validate_service_instance
 from paasta_tools.vitesscluster_tools import load_vitess_instance_config
 
+try:
+    from vault_tools.paasta_secret import get_client as get_vault_client
+    from vault_tools.paasta_secret import get_vault_url
+    from vault_tools.paasta_secret import get_vault_ca
+except ImportError:
+
+    def get_vault_client(url: str, capath: str) -> None:
+        pass
+
+    def get_vault_url(ecosystem: str) -> str:
+        return ""
+
+    def get_vault_ca(ecosystem: str) -> str:
+        return ""
+
+
 log = logging.getLogger(__name__)
 
 
@@ -1083,3 +1101,34 @@ def get_paasta_oapi_api_clustername(cluster: str, is_eks: bool) -> str:
     "eks-" prefix
     """
     return f"eks-{cluster}" if is_eks else cluster
+
+
+def get_current_ecosystem() -> str:
+    """Get current ecosystem from host configs, defaults to dev if no config is found"""
+    try:
+        with open("/nail/etc/ecosystem") as f:
+            return f.read().strip()
+    except IOError:
+        pass
+    return "devc"
+
+
+def get_service_auth_token() -> str:
+    """Uses instance profile to authenticate with Vault and generate token for service authentication"""
+    ecosystem = get_current_ecosystem()
+    vault_client = get_vault_client(get_vault_url(ecosystem), get_vault_ca(ecosystem))
+    vault_role = load_system_paasta_config().get_service_auth_vault_role()
+    metadata_provider = InstanceMetadataProvider(
+        iam_role_fetcher=InstanceMetadataFetcher(),
+    )
+    instance_credentials = metadata_provider.load().get_frozen_credentials()
+    vault_client.auth.aws.iam_login(
+        instance_credentials.access_key,
+        instance_credentials.secret_key,
+        instance_credentials.token,
+        mount_point="aws-iam",
+        role=vault_role,
+        use_token=True,
+    )
+    response = vault_client.secrets.identity.generate_signed_id_token(name=vault_role)
+    return response["data"]["token"]

paasta_tools/utils.py

@@ -2041,6 +2041,7 @@ class SystemPaastaConfigDict(TypedDict, total=False):
     secret_sync_delay_seconds: float
     use_multiple_log_readers: Optional[List[str]]
     service_auth_token_settings: ProjectedSAVolume
+    service_auth_vault_role: str
     always_authenticating_services: List[str]
     mysql_port_mappings: Dict
     vitess_images: Dict
@@ -2756,6 +2757,9 @@ def get_kube_clusters(self) -> Dict:
     def get_service_auth_token_volume_config(self) -> ProjectedSAVolume:
         return self.config_dict.get("service_auth_token_settings", {})
 
+    def get_service_auth_vault_role(self) -> str:
+        return self.config_dict.get("service_auth_vault_role", "service_authz")
+
     def get_always_authenticating_services(self) -> List[str]:
         return self.config_dict.get("always_authenticating_services", [])
 

tests/cli/test_cmds_local_run.py

@@ -394,6 +394,7 @@ def test_configure_and_run_command_uses_cmd_from_config(
     args.assume_role_arn = ""
     args.assume_pod_identity = False
     args.use_okta_role = False
+    args.use_service_auth_token = False
 
     mock_secret_provider_kwargs = {
         "vault_cluster_config": {},
@@ -436,6 +437,7 @@ def test_configure_and_run_command_uses_cmd_from_config(
         assume_pod_identity=False,
         assume_role_aws_account=None,
         use_okta_role=False,
+        use_service_auth_token=False,
     )
 
 
@@ -470,6 +472,7 @@ def test_configure_and_run_uses_bash_by_default_when_interactive(
     args.assume_role_arn = ""
     args.assume_pod_identity = False
     args.use_okta_role = False
+    args.use_service_auth_token = False
 
     return_code = configure_and_run_docker_container(
         docker_client=mock_docker_client,
@@ -511,6 +514,7 @@ def test_configure_and_run_uses_bash_by_default_when_interactive(
         assume_role_aws_account="dev",
         assume_pod_identity=False,
         use_okta_role=False,
+        use_service_auth_token=False,
     )
 
 
@@ -551,6 +555,7 @@ def test_configure_and_run_pulls_image_when_asked(
     args.assume_role_arn = ""
     args.assume_pod_identity = False
     args.use_okta_role = False
+    args.use_service_auth_token = False
 
     return_code = configure_and_run_docker_container(
         docker_client=mock_docker_client,
@@ -594,6 +599,7 @@ def test_configure_and_run_pulls_image_when_asked(
         assume_pod_identity=False,
         assume_role_aws_account="dev",
         use_okta_role=False,
+        use_service_auth_token=False,
     )
 
 
@@ -630,6 +636,7 @@ def test_configure_and_run_docker_container_defaults_to_interactive_instance(
         args.assume_role_arn = ""
         args.assume_pod_identity = False
         args.use_okta_role = False
+        args.use_service_auth_token = False
 
         mock_config = mock.create_autospec(AdhocJobConfig)
         mock_get_default_interactive_config.return_value = mock_config
@@ -673,6 +680,7 @@ def test_configure_and_run_docker_container_defaults_to_interactive_instance(
             assume_pod_identity=False,
             assume_role_aws_account="dev",
             use_okta_role=False,
+            use_service_auth_token=False,
         )
 
 

tests/cli/test_utils.py

@@ -24,6 +24,7 @@
 
 from paasta_tools.cli import utils
 from paasta_tools.cli.utils import extract_tags
+from paasta_tools.cli.utils import get_service_auth_token
 from paasta_tools.cli.utils import select_k8s_secret_namespace
 from paasta_tools.cli.utils import verify_instances
 from paasta_tools.kubernetes_tools import KubernetesDeploymentConfig
@@ -476,3 +477,50 @@ def test_select_k8s_secret_namespace():
 
     namespaces = {"a", "b"}
     assert select_k8s_secret_namespace(namespaces) in {"a", "b"}
+
+
+@patch("paasta_tools.cli.utils.load_system_paasta_config", autospec=True)
+@patch("paasta_tools.cli.utils.get_current_ecosystem", autospec=True)
+@patch("paasta_tools.cli.utils.InstanceMetadataProvider", autospec=True)
+@patch("paasta_tools.cli.utils.InstanceMetadataFetcher", autospec=True)
+@patch("paasta_tools.cli.utils.get_vault_client", autospec=True)
+@patch("paasta_tools.cli.utils.get_vault_url", autospec=True)
+@patch("paasta_tools.cli.utils.get_vault_ca", autospec=True)
+def test_get_service_auth_token(
+    mock_vault_ca,
+    mock_vault_url,
+    mock_get_vault_client,
+    mock_metadata_fetcher,
+    mock_metadata_provider,
+    mock_ecosystem,
+    mock_config,
+):
+    mock_ecosystem.return_value = "dev"
+    mock_config.return_value.get_service_auth_vault_role.return_value = "foobar"
+    mock_vault_client = mock_get_vault_client.return_value
+    mock_vault_client.secrets.identity.generate_signed_id_token.return_value = {
+        "data": {"token": "sometoken"},
+    }
+    assert get_service_auth_token() == "sometoken"
+    mock_instance_creds = (
+        mock_metadata_provider.return_value.load.return_value.get_frozen_credentials.return_value
+    )
+    mock_metadata_provider.assert_called_once_with(
+        iam_role_fetcher=mock_metadata_fetcher.return_value
+    )
+    mock_vault_url.assert_called_once_with("dev")
+    mock_vault_ca.assert_called_once_with("dev")
+    mock_get_vault_client.assert_called_once_with(
+        mock_vault_url.return_value, mock_vault_ca.return_value
+    )
+    mock_vault_client.auth.aws.iam_login.assert_called_once_with(
+        mock_instance_creds.access_key,
+        mock_instance_creds.secret_key,
+        mock_instance_creds.token,
+        mount_point="aws-iam",
+        role="foobar",
+        use_token=True,
+    )
+    mock_vault_client.secrets.identity.generate_signed_id_token.assert_called_once_with(
+        name="foobar"
+    )