From c459e086d1341b2a588c7fec4aae2a793d3b2bf1 Mon Sep 17 00:00:00 2001
From: Julien Lebot <julien.lebot@datadoghq.com>
Date: Wed, 31 Aug 2022 20:26:24 +0200
Subject: [PATCH] [Windows] Add a package invoke task for packaging the local
 builds of the custom action (#13245)

---
 .../agent/msi/localbuild/rebuild.bat          |  6 ++--
 tasks/customaction.py                         | 32 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/omnibus/resources/agent/msi/localbuild/rebuild.bat b/omnibus/resources/agent/msi/localbuild/rebuild.bat
index 65427b9ef2a69e..ba745b2c511d03 100644
--- a/omnibus/resources/agent/msi/localbuild/rebuild.bat
+++ b/omnibus/resources/agent/msi/localbuild/rebuild.bat
@@ -3,6 +3,8 @@ rem
 @set WD=%CD%
 cd %~dp0
 
+if NOT DEFINED OMNIBUS_BASE_DIR set OMNIBUS_BASE_DIR=C:\omnibus-ruby
+
 REM Copy few resource files locally, source.wxs AS IS and
 REM localization-en-us.wxl.erb replace template value currently via "Agent",
 REM unless localization-en-us.wxl already exhists (e.g. from previous copy)
@@ -35,7 +37,7 @@ heat.exe dir "c:\opt\datadog-agent" -nologo -srd -sreg -gg -cg ProjectDir -dr PR
 REM 
 REM run HEAT on the extras
 REM
-for /D %%D in (c:\omnibus-ruby\src\etc\datadog-agent\extra_package_files\*.*) do (
+for /D %%D in (%OMNIBUS_BASE_DIR%\src\etc\datadog-agent\extra_package_files\*.*) do (
     heat.exe dir %%D -nologo -srd -gg -cg Extra%%~nD -dr %%~nD -var "var.Extra%%~nD" -out "extra-%%~nD.wxs"
     set CANDLE_VARS=%CANDLE_VARS% -dExtra%%~nD="%%D"
     set WXS_LIST=%WX_LIST% extra-%%~nD.wxs
@@ -43,7 +45,7 @@ for /D %%D in (c:\omnibus-ruby\src\etc\datadog-agent\extra_package_files\*.*) do
 )
 
 @set wix_extension_switches=-ext WixUtilExtension
-candle -arch x64 %wix_extension_switches% -dProjectSourceDir="c:\opt\datadog-agent" -dExtraEXAMPLECONFSLOCATION="C:\omnibus-ruby\src\etc\datadog-agent\extra_package_files\EXAMPLECONFSLOCATION" project-files.wxs %WXS_LIST% source.wxs
+candle -arch x64 %wix_extension_switches% -dProjectSourceDir="c:\opt\datadog-agent" -dExtraEXAMPLECONFSLOCATION="%OMNIBUS_BASE_DIR%\src\etc\datadog-agent\extra_package_files\EXAMPLECONFSLOCATION" project-files.wxs %WXS_LIST% source.wxs
 
 if not "%ERRORLEVEL%" == "0" goto :done
 light -ext WixUIExtension -ext WixBalExtension %wix_extension_switches% -cultures:en-us -loc localization-en-us.wxl project-files.wixobj source.wixobj %WIXOBJ_LIST% -out ddagent.msi
diff --git a/tasks/customaction.py b/tasks/customaction.py
index 7f630212562c33..4bde31ee69074b 100644
--- a/tasks/customaction.py
+++ b/tasks/customaction.py
@@ -3,6 +3,7 @@
 """
 
 
+import glob
 import os
 import shutil
 import sys
@@ -102,3 +103,34 @@ def clean(_, arch="x64", debug=False):
         configuration = "Debug"
 
     shutil.rmtree(f"{CUSTOM_ACTION_ROOT_DIR}\\cal\\{arch}\\{configuration}", BIN_PATH)
+
+
+@task
+def package(
+    ctx,
+    vstudio_root=None,
+    omnibus_base_dir="c:\\omnibus-ruby",
+    arch="x64",
+    major_version='7',
+    debug=False,
+    rebuild=False,
+):
+    if os.getenv("OMNIBUS_BASE_DIR"):
+        omnibus_base_dir = os.getenv("OMNIBUS_BASE_DIR")
+    if rebuild:
+        clean(ctx, arch, debug)
+    build(ctx, vstudio_root, arch, major_version, debug)
+    for file in glob.glob(BIN_PATH + "\\customaction*"):
+        shutil.copy2(
+            file,
+            f"{omnibus_base_dir}\\src\\datadog-agent\\src\\github.com\\DataDog\\datadog-agent\\bin\\agent\\{os.path.basename(file)}",
+        )
+    cmd = "omnibus\\resources\\agent\\msi\\localbuild\\rebuild.bat"
+    res = ctx.run(cmd, warn=True)
+    if res.exited is None or res.exited > 0:
+        print(
+            color_message(
+                f"Failed to run \"{cmd}\"",
+                "orange",
+            )
+        )