Implement sanitizer handler for autorest modules (#1320)

* Implement sanitizer handler for autorest modules

* Generate code only when in Azure modules

powershell/cmdlets/class.ts

@@ -508,6 +508,11 @@ export class CmdletClass extends Class {
         }).toArray();
 
     this.NewImplementProcessRecordAsync();
+
+    if (this.state.project.azure) {
+      this.NewImplementWriteObject();
+    }
+
     this.debugMode = await this.state.getValue('debug', false);
 
     // json serialization
@@ -662,6 +667,14 @@ export class CmdletClass extends Class {
         if (!$this.state.project.azure) {
           yield $this.eventListener.syncSignal(Events.CmdletEndProcessing);
         }
+        else {
+          yield `var telemetryInfo = ${$this.state.project.serviceNamespace.moduleClass.declaration}.Instance.GetTelemetryInfo?.Invoke(__correlationId);`;
+          yield If('telemetryInfo != null', function* () {
+            yield 'telemetryInfo.TryGetValue("SanitizedProperties", out var sanitizedProperties);';
+            yield 'telemetryInfo.TryGetValue("InvocationName", out var invocationName);';
+            yield If('!string.IsNullOrEmpty(sanitizedProperties)', 'WriteWarning($"The output of cmdlet {invocationName ?? "Unknown"} may compromise security by showing the following secrets: {sanitizedProperties}. Learn more at https://go.microsoft.com/fwlink/?linkid=2258844");');
+          });
+        }
       });
 
     // debugging
@@ -900,6 +913,34 @@ export class CmdletClass extends Class {
     });
   }
 
+  private NewImplementWriteObject() {
+    const $this = this;
+    const sendToPipeline = new Parameter('sendToPipeline', dotnet.Object);
+    const enumerateCollection = new Parameter('enumerateCollection', dotnet.Bool);
+    const snglWriteObject = new Method('WriteObject', dotnet.Void, {
+      access: Access.Protected,
+      new: Modifier.New,
+      parameters: [sendToPipeline]
+    });
+    snglWriteObject.add(function* () {
+      yield `${$this.state.project.serviceNamespace.moduleClass.declaration}.Instance.SanitizeOutput?.Invoke(sendToPipeline, __correlationId);`;
+      yield 'base.WriteObject(sendToPipeline);';
+    });
+
+    const collWriteObject = new Method('WriteObject', dotnet.Void, {
+      access: Access.Protected,
+      new: Modifier.New,
+      parameters: [sendToPipeline, enumerateCollection]
+    });
+    collWriteObject.add(function* () {
+      yield `${$this.state.project.serviceNamespace.moduleClass.declaration}.Instance.SanitizeOutput?.Invoke(sendToPipeline, __correlationId);`;
+      yield 'base.WriteObject(sendToPipeline, enumerateCollection);';
+    });
+
+    $this.add(snglWriteObject);
+    $this.add(collWriteObject);
+  }
+
   private * ImplementCall(preProcess: PreProcess) {
     const $this = this;
     const operation = $this.operation;
@@ -2293,4 +2334,4 @@ export class CmdletClass extends Class {
       this.add(new Attribute(NotSuggestDefaultParameterSetAttribute));
     }
   }
-}
+}

powershell/generators/psm1.ts

@@ -26,12 +26,12 @@ export function getProfileExportScript(exportFolderScript: string, isAzure: bool
     # Load the last folder if no profile is selected
     $profileDirectory = $directories | Select-Object -Last 1
   }
-  
+
   if($profileDirectory) {
     Write-Information "Loaded Azure profile '$($profileDirectory.Name)' for module '$($instance.Name)'"
     $exportsPath = $profileDirectory.FullName
   }
-  
+
   if($exportsPath) {
     Get-ChildItem -Path $exportsPath -Recurse -Include '*.ps1' -File | ForEach-Object { . $_.FullName }
     $cmdletNames = Get-ScriptCmdlet -ScriptFolder $exportsPath
@@ -99,25 +99,31 @@ export async function generatePsm1(project: Project) {
 
   # Ask for the shared functionality table
   $VTable = Register-AzModule
-  
+
   # Tweaks the pipeline on module load
   $instance.OnModuleLoad = $VTable.OnModuleLoad
 
   # Following two delegates are added for telemetry
   $instance.GetTelemetryId = $VTable.GetTelemetryId
   $instance.Telemetry = $VTable.Telemetry
-  
+
+  # Delegate to sanitize the output object
+  $instance.SanitizeOutput = $VTable.SanitizerHandler
+
+  # Delegate to get the telemetry info
+  $instance.GetTelemetryInfo = $VTable.GetTelemetryInfo
+
 ${requestHandler}
-  
+
   # Gets shared parameter values
   $instance.GetParameterValue = $VTable.GetParameterValue
-  
+
   # Allows shared module to listen to events from this module
   $instance.EventListener = $VTable.EventListener
-  
+
   # Gets shared argument completers
   $instance.ArgumentCompleter = $VTable.ArgumentCompleter
-  
+
   # The name of the currently selected Azure profile
   $instance.ProfileName = $VTable.ProfileName
 `;
@@ -134,13 +140,13 @@ ${requestHandler}
 
   psm1.prepend('Generated', `
 ${azureInitialize}
- 
+
   # Load the custom module
   $customModulePath = Join-Path $PSScriptRoot '${project.psm1Custom}'
   if(Test-Path $customModulePath) {
     $null = Import-Module -Name $customModulePath
   }
-  
+
   # Export nothing to clear implicit exports
   Export-ModuleMember
 ${getProfileExportScript(`Join-Path $PSScriptRoot '${project.exportsFolder}'`, project.azure)}

powershell/module/module-class.ts

@@ -270,6 +270,15 @@ export class NewModuleClass extends Class {
       PSCmdlet, /* pscmdlet */
     )));
 
+    const sanitizerDelegate = new Alias('SanitizerDelegate', System.Action(
+      dotnet.Object, /* sendToPipeline */
+      dotnet.String  /* telemetryId */
+    ));
+    const getTelemetryInfoDelegate = new Alias('GetTelemetryInfoDelegate', System.Func(
+      dotnet.String /* telemetryId */,
+      /* returns */ System.Collections.Generic.Dictionary(System.String, System.String)
+    ));
+
     const tokenAudienceConverterDelegate = new Alias('TokenAudienceConverterDelegate',
       System.Func(
         dotnet.String,
@@ -288,6 +297,7 @@ export class NewModuleClass extends Class {
         tokenAudienceConverterDelegate.fullDefinition,
         System.Collections.Generic.IDictionary(dotnet.String, dotnet.Object)
       ));
+
     if (isDataPlane) {
       namespace.add(tokenAudienceConverterDelegate);
       namespace.add(authorizeRequestDelegate);
@@ -309,6 +319,10 @@ export class NewModuleClass extends Class {
     const AddAuthorizeRequestHandler = new Property('AddAuthorizeRequestHandler', authorizeRequestDelegate, { description: 'The delegate to call before each new request to add authorization.' });
     this.add(new Property('GetTelemetryId', getTelemetryIdDelegate, { description: 'The delegate to get the telemetry Id.' }));
     this.add(new Property('Telemetry', telemetryDelegate, { description: 'The delegate for creating a telemetry.' }));
+    const SanitizeOutput = this.add(new Property('SanitizeOutput', sanitizerDelegate, { description: 'The delegate to call in WriteObject to sanitize the output object.' }));
+    namespace.add(sanitizerDelegate);
+    const GetTelemetryInfo = this.add(new Property('GetTelemetryInfo', getTelemetryInfoDelegate, { description: 'The delegate to get the telemetry info.' }));
+    namespace.add(getTelemetryInfoDelegate);
     if (isDataPlane) {
       this.add(AddRequestUserAgentHandler);
       this.add(AddPatchRequestUriHandler);