From 9854cdce14b2c67aaa09dcfdcad83ac658434604 Mon Sep 17 00:00:00 2001
From: Chuck Driesler <cdriesler.iv@gmail.com>
Date: Thu, 2 May 2024 12:47:45 +0100
Subject: [PATCH] Feature: support automate function input secrets in C# SDK
 (#3324)

* introduce and use `SecretAttribute` during schema generation

* format

* format (for real)

* `pragma` and style

* Supressing CS8764 instead of CS8603

---------

Co-authored-by: Jedd Morgan <45512892+JR-Morgan@users.noreply.github.com>
---
 .../DataAnnotations/SecretAttribute.cs        |  7 ++++
 Automate/Speckle.Automate.Sdk/Runner.cs       | 34 +++++++++++++++++++
 2 files changed, 41 insertions(+)
 create mode 100644 Automate/Speckle.Automate.Sdk/DataAnnotations/SecretAttribute.cs

diff --git a/Automate/Speckle.Automate.Sdk/DataAnnotations/SecretAttribute.cs b/Automate/Speckle.Automate.Sdk/DataAnnotations/SecretAttribute.cs
new file mode 100644
index 0000000000..807e002928
--- /dev/null
+++ b/Automate/Speckle.Automate.Sdk/DataAnnotations/SecretAttribute.cs
@@ -0,0 +1,7 @@
+namespace Speckle.Automate.Sdk.DataAnnotations;
+
+/// <summary>
+/// If specified, the given function input will be redacted in all contexts.
+/// </summary>
+[AttributeUsage(AttributeTargets.All)]
+public sealed class SecretAttribute : Attribute { }
diff --git a/Automate/Speckle.Automate.Sdk/Runner.cs b/Automate/Speckle.Automate.Sdk/Runner.cs
index d3b0186adc..8a0d0b6496 100644
--- a/Automate/Speckle.Automate.Sdk/Runner.cs
+++ b/Automate/Speckle.Automate.Sdk/Runner.cs
@@ -1,8 +1,10 @@
 using System.CommandLine;
 using System.Diagnostics.CodeAnalysis;
+using Newtonsoft.Json;
 using Newtonsoft.Json.Schema;
 using Newtonsoft.Json.Schema.Generation;
 using Newtonsoft.Json.Serialization;
+using Speckle.Automate.Sdk.DataAnnotations;
 using Speckle.Automate.Sdk.Schema;
 using Speckle.Core.Logging;
 
@@ -131,6 +133,7 @@ public static async Task<int> Main<TInput>(string[] args, Func<AutomationContext
       (schemaFilePath) =>
       {
         JSchemaGenerator generator = new() { ContractResolver = new CamelCasePropertyNamesContractResolver() };
+        generator.GenerationProviders.Add(new SpeckleSecretProvider());
         JSchema schema = generator.Generate(typeof(TInput));
         schema.ToString(global::Newtonsoft.Json.Schema.SchemaVersion.Draft2019_09);
         File.WriteAllText(schemaFilePath, schema.ToString());
@@ -144,3 +147,34 @@ public static async Task<int> Main<TInput>(string[] args, Func<AutomationContext
     return returnCode;
   }
 }
+
+public class SpeckleSecretProvider : JSchemaGenerationProvider
+{
+  // `GetSchema` returning `null` indicates that the given type should not have a customised schema
+  // Nullability of JSchemaTypeGenerationContext appears to be incorrect.
+#pragma warning disable CS8764 // Nullability of return type doesn't match overridden member (possibly because of nullability attributes).
+  public override JSchema? GetSchema(JSchemaTypeGenerationContext context)
+  {
+    var attributes = context.MemberProperty?.AttributeProvider?.GetAttributes(false) ?? new List<Attribute>();
+    var isSecretString = attributes.Any(att => att is SecretAttribute);
+
+    if (isSecretString)
+    {
+      return CreateSchemaWithWriteOnly(context.ObjectType, context.Required);
+    }
+
+    return null;
+  }
+#pragma warning restore CS8764 // Nullability of return type doesn't match overridden member (possibly because of nullability attributes).
+
+
+  private JSchema CreateSchemaWithWriteOnly(Type type, Required required)
+  {
+    JSchemaGenerator generator = new();
+    JSchema schema = generator.Generate(type, required != Required.Always);
+
+    schema.WriteOnly = true;
+
+    return schema;
+  }
+}