Add --use_dotnet_agent_functions to deploy.py and Python service (#2292)

src/deployment/deploy.py

@@ -96,7 +96,12 @@ FUNC_TOOLS_ERROR = (
 )
 
 DOTNET_APPLICATION_SUFFIX = "-net"
-
+DOTNET_AGENT_FUNCTIONS = [
+    "agent_can_schedule",
+    "agent_commands",
+    "agent_events",
+    "agent_registration",
+]
 logger = logging.getLogger("deploy")
 
 
@@ -156,6 +161,7 @@ class Client:
         admins: List[UUID],
         allowed_aad_tenants: List[UUID],
         enable_dotnet: List[str],
+        use_dotnet_agent_functions: bool,
     ):
         self.subscription_id = subscription_id
         self.resource_group = resource_group
@@ -191,6 +197,7 @@ class Client:
         self.arm_template = bicep_to_arm(bicep_template)
 
         self.enable_dotnet = enable_dotnet
+        self.use_dotnet_agent_functions = use_dotnet_agent_functions
 
         machine = platform.machine()
         system = platform.system()
@@ -618,6 +625,7 @@ class Client:
             "signedExpiry": {"value": expiry},
             "multi_tenant_domain": multi_tenant_domain,
             "workbookData": {"value": self.workbook_data},
+            "use_dotnet_agent_functions": {"value": self.use_dotnet_agent_functions},
         }
         deployment = Deployment(
             properties=DeploymentProperties(
@@ -1099,12 +1107,7 @@ class Client:
             def expand_agent(f: str) -> List[str]:
                 # 'agent' is permitted as a shortcut for the agent functions
                 if f == "agent":
-                    return [
-                        "agent_can_schedule",
-                        "agent_commands",
-                        "agent_events",
-                        "agent_registration",
-                    ]
+                    return DOTNET_AGENT_FUNCTIONS
                 else:
                     return [f]
 
@@ -1363,6 +1366,11 @@ def main() -> None:
         "their functions and enable corresponding dotnet functions in the Azure "
         "Function App deployment",
     )
+    parser.add_argument(
+        "--use_dotnet_agent_functions",
+        action="store_true",
+        help="Tell the OneFuzz agent to use the dotnet endpoint",
+    )
     args = parser.parse_args()
 
     if shutil.which("func") is None:
@@ -1393,6 +1401,7 @@ def main() -> None:
         admins=args.set_admins,
         allowed_aad_tenants=args.allowed_aad_tenants or [],
         enable_dotnet=args.enable_dotnet,
+        use_dotnet_agent_functions=args.use_dotnet_agent_functions,
     )
     if args.verbose:
         level = logging.DEBUG
@@ -1403,6 +1412,17 @@ def main() -> None:
 
     logging.getLogger("deploy").setLevel(logging.INFO)
 
+    if args.use_dotnet_agent_functions:
+        # validate that the agent functions are actually enabled
+        if not (
+            "agent" in args.enable_dotnet
+            or all(map(lambda f: f in args.enable_dotnet, DOTNET_AGENT_FUNCTIONS))
+        ):
+            logger.error(
+                "If --use_dotnet_agent_functions is set, all agent functions must be enabled (--enable_dotnet agent)."
+            )
+            sys.exit(1)
+
     if args.rbac_only:
         logger.warning(
             "'rbac_only' specified. The deployment will execute "