From 5f0c6b326f82438570468578951c0d17a719169c Mon Sep 17 00:00:00 2001
From: Olga Naidjonoka <olga_naidjonoka@epam.com>
Date: Thu, 3 Oct 2024 13:51:23 +0300
Subject: [PATCH] added unzip

---
 .buildkite/scripts/agentbeat/prepare_env.py | 100 ++++++++++++++++----
 1 file changed, 81 insertions(+), 19 deletions(-)

diff --git a/.buildkite/scripts/agentbeat/prepare_env.py b/.buildkite/scripts/agentbeat/prepare_env.py
index db42e64ee8fc..9da84ec19f9f 100755
--- a/.buildkite/scripts/agentbeat/prepare_env.py
+++ b/.buildkite/scripts/agentbeat/prepare_env.py
@@ -3,6 +3,14 @@
 import platform
 import subprocess
 import sys
+import tarfile
+
+PATH = 'x-pack/agentbeat/build/distributions'
+
+
+def log(msg):
+    sys.stdout.write(f'{msg}\n')
+    sys.stdout.flush()
 
 
 def get_os() -> str:
@@ -12,56 +20,110 @@ def get_os() -> str:
 def get_arch() -> str:
     arch = platform.machine().lower()
 
-    if arch == "amd64":
-        return "x86_64"
+    if arch == 'amd64':
+        return 'x86_64'
     else:
         return arch
 
 
 def get_artifact_extension(agent_os) -> str:
-    if agent_os == "windows":
-        return "zip"
+    if agent_os == 'windows':
+        return 'zip'
     else:
-        return "tar.gz"
+        return 'tar.gz'
 
 
-def download_agentbeat_artifact(agent_os, agent_arch):
-    print(" ")
+def get_artifact_pattern() -> str:
+    agent_os = get_os()
+    agent_arch = get_arch()
     extension = get_artifact_extension(agent_os)
-    pattern = f"x-pack/agentbeat/build/distributions/agentbeat-*-{agent_os}-{agent_arch}.{extension}"
+    return f'{PATH}/agentbeat-*-{agent_os}-{agent_arch}.{extension}'
+
 
+def download_agentbeat(pattern, path) -> str:
+    log('--- Downloading agentbeat')
     try:
         subprocess.run(
-            ["buildkite-agent", "artifact", "download", pattern, ".",
-             "--build", "01924d2b-b061-45ae-a106-e885584ff26f",
-             "--step", "agentbeat-package-linux"],
+            ['buildkite-agent', 'artifact', 'download', pattern, '.',
+             '--build', '01924d2b-b061-45ae-a106-e885584ff26f',
+             '--step', 'agentbeat-package-linux'],
             check=True, stdout=sys.stdout, stderr=sys.stderr, text=True)
     except subprocess.CalledProcessError:
         exit(1)
 
+    return get_filename(path)
+
+
+def get_filename(path) -> str:
+    try:
+        out = subprocess.run(
+            ['ls', '-p', path],
+            check=True, capture_output=True, text=True)
+        return out.stdout.strip()
+    except subprocess.CalledProcessError:
+        exit(1)
+
+
+def extract_agentbeat(filename):
+    log('--- Extracting agentbeat')
+    filepath = PATH + '/' + filename
+
+
+    if filepath.endswith('.zip'):
+        unzip_agentbeat(filepath)
+    else:
+        untar_agentbeat(filepath)
+    log('Successfully extracted agentbeat')
+
 
-def unzip_agentbeat():
-    print("todo unzip")
+def unzip_agentbeat(filepath):
     try:
         subprocess.run(
-            ["unzip"],
+            ['unzip', filepath],
             check=True, stdout=sys.stdout, stderr=sys.stderr, text=True)
+    except subprocess.CalledProcessError as e:
+        sys.stderr.write(str(e))
+        exit(1)
+
+
+def untar_agentbeat(filepath):
+    try:
+        with tarfile.open(filepath, 'r:gz') as tar:
+            tar.list()
+            tar.extractall()
+    except Exception as e:
+        sys.stderr.write(str(e))
+        exit(1)
+
+
+def add_to_path(filepath):
+    try:
+        out = subprocess.run(
+            ['ls', '-p', 'x-pack/agentbeat/build/distributions'],
+            check=True, capture_output=True, text=True)
+
     except subprocess.CalledProcessError:
         exit(1)
 
+    print(filepath)
+
 
 def install_synthetics():
-    print("--- Installing @elastic/synthetics")
+    log('--- Installing @elastic/synthetics')
 
     try:
         subprocess.run(
-            ["npm install -g @elastic/synthetics"],
+            ['npm install -g @elastic/synthetics'],
             check=True
         )
     except subprocess.CalledProcessError:
-        print("Failed to install @elastic/synthetics")
+        log('Failed to install @elastic/synthetics')
         exit(1)
 
-print("--- OS Data: " + get_os() + " " + get_arch())
-download_agentbeat_artifact(get_os(), get_arch())
+print('OS Data: ' + get_os() + ' ' + get_arch())
+
+artifact_pattern = get_artifact_pattern()
+archive = download_agentbeat(artifact_pattern, PATH)
+extract_agentbeat(archive)
+add_to_path(archive)
 # install_synthetics()