diff --git a/src/acat/azext_acat/custom.py b/src/acat/azext_acat/custom.py index 6fefabfd0b..17161e22b5 100644 --- a/src/acat/azext_acat/custom.py +++ b/src/acat/azext_acat/custom.py @@ -9,37 +9,17 @@ # pylint: disable=too-many-statements from knack.log import get_logger -import datetime -from azure.cli.core._profile import Profile -from azure.cli.core.aaz._client import AAZMgmtClient from .aaz.latest.app_compliance_automation import Onboard as _AcatOnboard from .aaz.latest.app_compliance_automation.report import List as _AcatListReport from .aaz.latest.app_compliance_automation.report import Show as _AcatShowReport from .aaz.latest.app_compliance_automation.report import Create as _AcatCreateReport from .aaz.latest.app_compliance_automation.report import Update as _AcatUpdateReport from .aaz.latest.app_compliance_automation.report import Delete as _AcatDeleteReport -from azure.cli.core.aaz import register_client,AAZClientConfiguration from azure.cli.core.commands import LongRunningOperation - +from .utils import * logger = get_logger(__name__) -@register_client("DupAadTokenClient") -class AAZDupAadTokenClient(AAZMgmtClient): - @classmethod - def _build_configuration(cls, ctx, credential, **kwargs): - from azure.cli.core.auth.util import resource_to_scopes - from azure.core.pipeline.policies import HeadersPolicy - - token="Bearer "+credential.get_token().token - headers_policy=HeadersPolicy(**kwargs) - headers_policy.add_header("x-ms-aad-user-token",token) - kwargs["headers_policy"]=headers_policy - - return AAZClientConfiguration( - credential=credential, - credential_scopes=resource_to_scopes(ctx.cli_ctx.cloud.endpoints.active_directory_resource_id), - **kwargs - ) + class OnboardAcat(_AcatOnboard): class OnboardAcatWithDupAadToken(_AcatOnboard.ProviderActionsOnboard): @@ -51,7 +31,6 @@ def _execute_operations(self): self.post_operations() - class ListAcatReport(_AcatListReport): class ListAcatReportWithDupAadToken(_AcatListReport.ReportList): CLIENT_TYPE = "DupAadTokenClient" @@ -134,8 +113,9 @@ def pre_operations(self): cmd.cli_ctx, 'Starting {}'.format(cmd.name))(poller) args = self.ctx.args - args.time_zone=args.nullable_time_zone_with_default + args.time_zone=args.time_zone_with_default args.trigger_time=args.trigger_time_with_default + class UpdateAcatReport(_AcatUpdateReport): class UpdateAcatReportWithDupAadToken(_AcatUpdateReport.ReportUpdate): @@ -146,6 +126,7 @@ def _execute_operations(self): yield self.UpdateAcatReportWithDupAadToken(ctx=self.ctx)() self.post_operations() + class DeleteAcatReport(_AcatDeleteReport): class DeleteAcatReportWithDupAadToken(_AcatDeleteReport.ReportDelete): CLIENT_TYPE = "DupAadTokenClient" @@ -157,19 +138,3 @@ def _execute_operations(self): # utils -def ParseSubsFromResources(resourceList): - """Parse subscriptions from resources""" - subs = set() - for resource in resourceList: - sub = resource.resource_id._data.split("/")[2] - subs.add(sub) - return list(subs) - -def GetLocalTimeZone(): - """Get local time zone""" - return str(datetime.datetime.now(datetime.timezone.utc).astimezone().tzinfo) - -def GetClosestFullHour(): - """Get closest full hour""" - now = datetime.datetime.now() - return str(now.replace(minute=0, second=0, microsecond=0) + datetime.timedelta(hours=1)) diff --git a/src/acat/azext_acat/utils.py b/src/acat/azext_acat/utils.py new file mode 100644 index 0000000000..7b3e7ffd61 --- /dev/null +++ b/src/acat/azext_acat/utils.py @@ -0,0 +1,38 @@ +from azure.cli.core.aaz import register_client,AAZClientConfiguration +from azure.cli.core.aaz._client import AAZMgmtClient +import datetime + +@register_client("DupAadTokenClient") +class AAZDupAadTokenClient(AAZMgmtClient): + @classmethod + def _build_configuration(cls, ctx, credential, **kwargs): + from azure.cli.core.auth.util import resource_to_scopes + from azure.core.pipeline.policies import HeadersPolicy + + token="Bearer "+credential.get_token().token + headers_policy=HeadersPolicy(**kwargs) + headers_policy.add_header("x-ms-aad-user-token",token) + kwargs["headers_policy"]=headers_policy + + return AAZClientConfiguration( + credential=credential, + credential_scopes=resource_to_scopes(ctx.cli_ctx.cloud.endpoints.active_directory_resource_id), + **kwargs + ) + +def ParseSubsFromResources(resourceList): + """Parse subscriptions from resources""" + subs = set() + for resource in resourceList: + sub = resource.resource_id._data.split("/")[2] + subs.add(sub) + return list(subs) + +def GetLocalTimeZone(): + """Get local time zone""" + return str(datetime.datetime.now(datetime.timezone.utc).astimezone().tzinfo) + +def GetClosestFullHour(): + """Get closest full hour""" + now = datetime.datetime.now() + return str(now.replace(minute=0, second=0, microsecond=0) + datetime.timedelta(hours=1))