From e5285f2f78c28baa1881a81d300c1fba4242adf8 Mon Sep 17 00:00:00 2001
From: vakarisz <vakarisz@yahoo.com>
Date: Tue, 10 May 2022 16:57:36 +0300
Subject: [PATCH] Island: Add custom error and sketch out AWS command results

---
 .../cc/services/aws/aws_command_runner.py     | 27 ++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/monkey/monkey_island/cc/services/aws/aws_command_runner.py b/monkey/monkey_island/cc/services/aws/aws_command_runner.py
index 78dfdc21c..b6e9aa16e 100644
--- a/monkey/monkey_island/cc/services/aws/aws_command_runner.py
+++ b/monkey/monkey_island/cc/services/aws/aws_command_runner.py
@@ -85,6 +85,10 @@ def _run_command_async(
     return command_id
 
 
+class AWSCommandError(Exception):
+    pass
+
+
 def _wait_for_command_to_complete(
     aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str
 ):
@@ -94,9 +98,11 @@ def _wait_for_command_to_complete(
     while not timer.is_expired():
         time.sleep(STATUS_CHECK_SLEEP_TIME)
 
-        command_status = aws_client.get_command_invocation(
+        command_result = aws_client.get_command_invocation(
             CommandId=command_id, InstanceId=target_instance_id
-        )["Status"]
+        )
+        command_status = command_result["Status"]
+
         logger.debug(f"Command {command_id} status: {command_status}")
 
         if command_status == "Success":
@@ -104,4 +110,19 @@ def _wait_for_command_to_complete(
 
         if command_status != "InProgress":
             # TODO: Create an exception for this occasion and raise it with useful information.
-            raise Exception("COMMAND FAILED")
+            raise AWSCommandError(
+                f"AWS command failed." f" Command invocation contents: {command_result}"
+            )
+
+
+def _fetch_command_results(
+    aws_client: botocore.client.BaseClient, target_instance_id: str, command_id: str
+):
+    command_results = aws_client.ssm.get_command_invocation(
+        CommandId=command_id, InstanceId=target_instance_id
+    )
+    # TODO: put these into a dataclass and return
+    # self.is_successful(command_info, True)
+    # command_results["ResponseCode"]
+    # command_results["StandardOutputContent"]
+    # command_results["StandardErrorContent"]