From e524f625f797682c996dfa46214a431902f849ae Mon Sep 17 00:00:00 2001
From: YanaXu <sheshouxy@gmail.com>
Date: Fri, 13 Sep 2024 15:34:50 +0800
Subject: [PATCH] move pre cleanup to main and add pre-if and post-if

---
 action.yml  | 2 +-
 src/main.ts | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/action.yml b/action.yml
index af6c2739b..44c1f66a6 100644
--- a/action.yml
+++ b/action.yml
@@ -39,6 +39,6 @@ branding:
   color: 'blue'
 runs:
   using: 'node20'
-  pre: 'lib/cleanup/index.js'
   main: 'lib/main/index.js'
+  post-if: (!env.AZURE_LOGIN_POST_CLEANUP || env.AZURE_LOGIN_POST_CLEANUP != 'false')
   post: 'lib/cleanup/index.js'
diff --git a/src/main.ts b/src/main.ts
index b3be983a7..add837a9e 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -1,5 +1,5 @@
 import * as core from '@actions/core';
-import { setUserAgent } from './common/Utils'; 
+import { cleanupAzCLIAccounts, cleanupAzPSAccounts, setUserAgent } from './common/Utils';
 import { AzPSLogin } from './PowerShell/AzPSLogin';
 import { LoginConfig } from './common/LoginConfig';
 import { AzureCliLogin } from './Cli/AzureCliLogin';
@@ -7,6 +7,13 @@ import { AzureCliLogin } from './Cli/AzureCliLogin';
 async function main() {
     try {
         setUserAgent();
+        const preCleanup: string = process.env.AZURE_LOGIN_PRE_CLEANUP;
+        if ('true' == preCleanup) {
+            await cleanupAzCLIAccounts();
+            if (core.getInput('enable-AzPSSession').toLowerCase() === "true") {
+                await cleanupAzPSAccounts();
+            }
+        }
 
         // prepare the login configuration
         var loginConfig = new LoginConfig();