update apply script; add failsafe in case SSH connection times out

runprism · Sep 13, 2023 · 4590555 · 4590555
1 parent 514a85c
commit 4590555
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 8 deletions.
diff --git a/prism/agents/ec2.py b/prism/agents/ec2.py
@@ -451,7 +451,7 @@ def add_ingress_rule(self,
                     'IpProtocol': 'tcp',
                     'FromPort': 22,
                     'ToPort': 22,
-                    'IpRanges': [{'CidrIp': curr_cidr}]
+                    'IpRanges': [curr_cidr]
                 },
             ]
         else:
@@ -693,7 +693,7 @@ def _create_exception(resource):
                     KeyName=instance_name,
                     MinCount=1,
                     MaxCount=1,
-                    ImageId="ami-0889a44b331db0194",
+                    ImageId="ami-01c647eace872fc02",
                     TagSpecifications=[
                         {
                             'ResourceType': 'instance',
@@ -1071,7 +1071,6 @@ def apply(self):
                 # If the current IP address is already whitelisted, then just add
                 # 0.0.0.0/0 to the ingress rules.
                 if added_ip is None:
-                    self.args.whitelist_all = True
                     prism.prism_logging.DEFAULT_LOGGER.agent(  # type: ignore
                         f"{prism.ui.AGENT_EVENT}{self.instance_name}{prism.ui.AGENT_WHICH_BUILD}[build]{prism.ui.RESET}  | Current IP address already whitelisted...whitelisting 0.0.0.0/0"  # noqa: E501
                     )

diff --git a/prism/agents/scripts/apply.sh b/prism/agents/scripts/apply.sh
@@ -123,7 +123,7 @@ if [ $exit_code -eq 1 ]; then
 fi
 
 # Move all folders into the root folder
-ssh -i ${pem_path} ${user}@${public_dns_name} 'cd ~ && for dir in */; do sudo rm -rf ../../$dir; do sudo mv $dir ../../$dir ; done'
+ssh -i ${pem_path} ${user}@${public_dns_name} 'cd ~ && for dir in */; do sudo rm -rf ../../$dir; sudo mv -f $dir ../../ ; done'
 exit_code=$?
 if [ $exit_code -eq 1 ]; then
 	exit 1

diff --git a/prism/cli/agent.py b/prism/cli/agent.py
@@ -31,7 +31,7 @@
 import json
 from pathlib import Path
 import os
-from typing import List
+from typing import List, Union
 
 
 ####################
@@ -73,7 +73,7 @@ def build_agent(self,
         agent_yml_path: Path,
         agent_filename: str,
         prism_project: PrismProject,
-    ):
+    ) -> Union[int, Agent]:
         """
         Create the agent instance
         """
@@ -116,8 +116,11 @@ def build_agent(self,
         # run the agent (`prism agent build`). We run the agent in a separate function,
         # so focus on just building it for now.
         else:
-            agent.apply()
-            return agent
+            returncode = agent.apply()
+            if returncode["return_code"] != 0:
+                return 0
+            else:
+                return agent
 
     def run_agent(self, agent: Agent):
         """