Extend generate token

azext_iot/_params.py

@@ -232,6 +232,13 @@ def load_arguments(self, _):
             arg_type=get_three_state_flag(),
             help="Flag indicating edge enablement.",
         )
+        context.argument(
+            "connection_string",
+            options_list=["--connection-string", "--cs"],
+            help="Target connection string. This bypasses the IoT Hub registry and generates the SAS token directly"
+            " from the supplied symmetric key without further validation. All other command parameters aside from"
+            " duration will be ignored."
+        )
 
     with self.argument_context("iot hub") as context:
         context.argument(

azext_iot/common/shared.py

@@ -267,3 +267,20 @@ class IoTHubStateType(Enum):
     KeyEncryptionKeyRevoking = "KeyEncryptionKeyRevoking"
     KeyEncryptionKeyRevoked = "KeyEncryptionKeyRevoked"
     ReActivating = "ReActivating"
+
+
+class ConnectionStringParser(Enum):
+    """
+        All connection string parser with respective functions
+    """
+    from azext_iot.common._azure import (
+        parse_iot_device_connection_string,
+        parse_iot_device_module_connection_string,
+        parse_pnp_connection_string,
+        parse_iot_hub_connection_string
+    )
+
+    Module = parse_iot_device_module_connection_string
+    Device = parse_iot_device_connection_string
+    PnP = parse_pnp_connection_string
+    IotHub = parse_iot_hub_connection_string

azext_iot/operations/hub.py

@@ -29,6 +29,7 @@
     RenewKeyType,
     IoTHubStateType,
     DeviceAuthApiType,
+    ConnectionStringParser,
 )
 from azext_iot.iothub.providers.discovery import IotHubDiscovery
 from azext_iot.common.utility import (
@@ -1907,6 +1908,7 @@ def iot_get_sas_token(
     login=None,
     module_id=None,
     auth_type_dataplane=None,
+    connection_string=None,
 ):
     key_type = key_type.lower()
     policy_name = policy_name.lower()
@@ -1924,6 +1926,14 @@ def iot_get_sas_token(
             "You are unable to get sas token for module without device information."
         )
 
+    if connection_string:
+        return {
+            DeviceAuthApiType.sas.value: _iot_build_sas_token_from_cs(
+                connection_string,
+                duration,
+            ).generate_sas_token()
+        }
+
     return {
         DeviceAuthApiType.sas.value: _iot_build_sas_token(
             cmd,
@@ -1940,6 +1950,45 @@ def iot_get_sas_token(
     }
 
 
+def _iot_build_sas_token_from_cs(connection_string, duration=3600):
+    uri = None
+    policy = None
+    key = None
+
+    parsed_cs = None
+    all_parsers = [
+        ConnectionStringParser.Module,
+        ConnectionStringParser.Device,
+        ConnectionStringParser.PnP,
+        ConnectionStringParser.IotHub,
+    ]
+
+    for parser in all_parsers:
+        try:
+            parsed_cs = parser(connection_string)
+
+            if "SharedAccessKeyName" in parsed_cs:
+                policy = parsed_cs["SharedAccessKeyName"]
+            key = parsed_cs["SharedAccessKey"]
+
+            if parser in [ConnectionStringParser.IotHub, ConnectionStringParser.PnP]:
+                uri = parsed_cs["HostName"]
+            elif parser == ConnectionStringParser.Module:
+                uri = "{}/devices/{}/modules/{}".format(
+                    parsed_cs["HostName"], parsed_cs["DeviceId"], parsed_cs["ModuleId"]
+                )
+            elif parser == ConnectionStringParser.Device:
+                uri = "{}/devices/{}".format(parsed_cs["HostName"], parsed_cs["DeviceId"])
+            else:
+                raise CLIError("Connection String did not match any presets")
+
+            return SasTokenAuthentication(uri, policy, key, duration)
+        except ValueError:
+            continue
+
+    raise CLIError("Connection String did not match any presets")
+
+
 def _iot_build_sas_token(
     cmd,
     hub_name=None,
@@ -1972,6 +2021,7 @@ def _iot_build_sas_token(
     uri = None
     policy = None
     key = None
+
     if device_id:
         logger.info(
             'Obtaining device "%s" details from registry, using IoT Hub policy "%s"',

azext_iot/tests/iothub/test_iothub_utilities_unit.py

@@ -0,0 +1,87 @@
+# coding=utf-8
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azext_iot.common.sas_token_auth import SasTokenAuthentication
+import pytest
+from knack.cli import CLIError
+from azext_iot.operations import hub as subject
+from azext_iot.tests.generators import generate_generic_id
+
+
+def generate_valid_cs(validate_pairs=[]):
+    host_name = generate_generic_id()
+    shared_access_key = generate_generic_id()
+    cs = f"HostName={host_name};"
+    input_pairs = dict((k, generate_generic_id()) for k in validate_pairs)
+    policy = input_pairs["SharedAccessKeyName"] if "SharedAccessKeyName" in input_pairs else None
+
+    for key, value in input_pairs.items():
+        cs += "{}={};".format(
+            key, value
+        )
+
+    cs = f"{cs}SharedAccessKey={shared_access_key}"
+    uri = host_name
+    if "DeviceId" in input_pairs:
+        uri = f"{uri}/devices/{input_pairs['DeviceId']}"
+    if "ModuleId" in input_pairs:
+        uri = f"{uri}/modules/{input_pairs['ModuleId']}"
+
+    return {
+        "connection_string": cs,
+        "uri": uri,
+        "policy": policy,
+        "key": shared_access_key
+    }
+
+
+class TestGenerateSasToken:
+    @pytest.mark.parametrize(
+        "duration, req",
+        [
+            (3600, generate_valid_cs(["DeviceId"])),
+            (30, generate_valid_cs(["DeviceId"])),
+            (60000, generate_valid_cs(["DeviceId"])),
+            (3600, generate_valid_cs(["RepositoryId", "SharedAccessKeyName"])),
+            (3600, generate_valid_cs(["SharedAccessKeyName"])),
+            (3600, generate_valid_cs(["DeviceId"])),
+            (3600, generate_valid_cs(["DeviceId", "ModuleId"])),
+            (3600, generate_valid_cs(["Test", "DeviceId", "ModuleId"])),
+            (3600, generate_valid_cs(["RepositoryId", "DeviceId", "ModuleId"])),
+        ],
+    )
+    def test_generate_sas_token_from_cs(self, mocker, fixture_cmd, duration, req):
+        patched_time = mocker.patch(
+            "azext_iot.common.sas_token_auth.time"
+        )
+        patched_time.return_value = 0
+        result = subject.iot_get_sas_token(
+            cmd=fixture_cmd,
+            connection_string=req["connection_string"],
+            duration=duration
+        )
+
+        duration = duration if duration else 3600
+        expected_sas = SasTokenAuthentication(
+            req["uri"], req["policy"], req["key"], duration
+        ).generate_sas_token()
+        assert result["sas"] == expected_sas
+
+    @pytest.mark.parametrize(
+        "req",
+        [
+            (generate_valid_cs()),
+            (generate_valid_cs(["ModuleId"])),
+            (generate_valid_cs(["RepositoryId"])),
+            (generate_valid_cs(["Test"]))
+        ],
+    )
+    def test_generate_sas_token_from_cs_error(self, mocker, fixture_cmd, req):
+        with pytest.raises(CLIError):
+            subject.iot_get_sas_token(
+                cmd=fixture_cmd,
+                connection_string=req["connection_string"],
+            )