Merge pull request #34 from runprism/ipv6-bug

Issue #32: SSH into an EC2 instance using an IPv6 address

prism/agents/ec2.py

@@ -443,12 +443,15 @@ def add_ingress_rule(self,
 
         # Add rule
         if ip_address_type == IpAddressType('ipv4'):
+            curr_cidr = {'CidrIp': f'{external_ip}/32'}
+            if external_ip == "0.0.0.0":
+                curr_cidr = {'CidrIp': '0.0.0.0/0'}
             ip_permissions = [
                 {
                     'IpProtocol': 'tcp',
                     'FromPort': 22,
                     'ToPort': 22,
-                    'IpRanges': [{'CidrIp': f'{external_ip}/32'}]
+                    'IpRanges': [curr_cidr]
                 },
             ]
         else:
@@ -457,14 +460,22 @@ def add_ingress_rule(self,
                     'IpProtocol': 'tcp',
                     'FromPort': 22,
                     'ToPort': 22,
-                    'Ipv6Ranges': [{'CidrIpv6': f'{external_ip}/128'}]
+                    'Ipv6Ranges': [{'CidrIpv6': '::/0'}]
                 },
             ]
 
-        _ = ec2_client.authorize_security_group_ingress(
-            GroupId=security_group_id,
-            IpPermissions=ip_permissions
-        )
+        try:
+            _ = ec2_client.authorize_security_group_ingress(
+                GroupId=security_group_id,
+                IpPermissions=ip_permissions
+            )
+            return external_ip
+
+        # If the rule already exists, then we're all gucci
+        except botocore.exceptions.ClientError as e:
+            if "the specified rule" in str(e) and "already exists" in str(e):
+                return None
+            raise e
 
     def create_new_security_group(self,
         ec2_client: Any,
@@ -561,11 +572,12 @@ def check_ingress_ip(self,
 
         # If SSH traffic from the current IP address it not allowed, then authorize it.
         if not ip_allowed:
-            self.add_ingress_rule(
+            return self.add_ingress_rule(
                 ec2_client,
                 security_group_id,
                 external_ip
             )
+        return None
 
     def create_instance(self,
         ec2_client: Any,
@@ -681,7 +693,7 @@ def _create_exception(resource):
                     KeyName=instance_name,
                     MinCount=1,
                     MaxCount=1,
-                    ImageId="ami-0889a44b331db0194",
+                    ImageId="ami-01c647eace872fc02",
                     TagSpecifications=[
                         {
                             'ResourceType': 'instance',
@@ -1020,7 +1032,7 @@ def apply(self):
 
             # Build the shell command
             cmd = [
-                '/bin/sh', self.AGENT_APPLY_SCRIPT,
+                '/bin/bash', self.AGENT_APPLY_SCRIPT,
                 '-r', str(requirements_txt_path),
                 '-p', str(pem_key_path),
                 '-u', user,
@@ -1041,6 +1053,38 @@ def apply(self):
                     f"{prism.ui.AGENT_EVENT}{self.instance_name}{prism.ui.AGENT_WHICH_BUILD}[build]{prism.ui.RESET} | {line.rstrip()}"  # noqa: E501
                 )
 
+            # A return code of 8 indicates that the SSH connection timed out. Try
+            # whitelisting the current IP and try again.
+            if returncode == 8:
+                # For mypy
+                if self.security_group_id is None:
+                    raise ValueError("`security_group_id` is still None!")
+                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}  | SSH connection timed out...checking security group ingress rules and trying again"  # noqa: E501
+                )
+
+                # Add the current IP to the security ingress rules
+                added_ip = self.check_ingress_ip(
+                    self.ec2_client, self.security_group_id
+                )
+
+                # 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:
+                    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
+                    )
+                    self.add_ingress_rule(
+                        self.ec2_client,
+                        self.security_group_id,
+                        "0.0.0.0",
+                    )
+
+                # Try again
+                _, err, returncode = self.stream_logs(
+                    cmd, prism.ui.AGENT_WHICH_BUILD, "build"
+                )
+
             # If the return code is non-zero, then an error occurred. Delete all of the
             # resources so that the user can try again.
             if returncode != 0:
@@ -1075,7 +1119,7 @@ def run(self):
 
         # The agent data should exist...Build the shell command
         cmd = [
-            '/bin/sh', self.AGENT_RUN_SCRIPT,
+            '/bin/bash', self.AGENT_RUN_SCRIPT,
             '-p', str(self.pem_key_path),
             '-u', 'ec2-user',
             '-n', self.public_dns_name,

prism/agents/scripts/apply.sh

@@ -16,13 +16,18 @@ done
 # ssh into the project so that we can authenticate our key
 while true
 do
-  	ssh -o "StrictHostKeyChecking no" -i "${pem_path}" "${user}@${public_dns_name}" exit 2>/dev/null 2>&1
-    if [ $? -eq 0 ]; then
+  	errormessage=`ssh -o "StrictHostKeyChecking no" -i "${pem_path}" "${user}@${public_dns_name}" exit 2>/dev/null 2>&1`
+    if [ -z "$errormessage" ]; then
         echo "SSH connection succeeded!"
         break
     else
-        echo "SSH connection failed. Retrying in 5 seconds..."
-        sleep 5
+		if [[ "$errormessage" =~ "Operation timed out" ]]; then
+			echo "SSH connection failed."
+			exit 8
+		else
+        	echo "SSH connection refused. Retrying in 5 seconds..."
+        	sleep 5
+		fi
     fi
 done
 
@@ -52,20 +57,32 @@ source ~/.venv/${project_name}/bin/activate
 pip install --upgrade pip
 pip install -r requirements.txt
 EOF
+	exit_code=$?
+	if [ $exit_code -eq 1 ]; then
+		exit 1
+	fi
 fi
 
 # Log
 echo "Updating remote project and file paths"
 
 # Copy project directory and other copy paths into the EC2 instance
 ssh -i ${pem_path} ${user}@${public_dns_name} "sudo mkdir -p .${project_dir}; sudo chmod 777 -R .${project_dir}"
+exit_code=$?
+if [ $exit_code -eq 1 ]; then
+	exit 1
+fi
 scp -r -i ${pem_path} ${project_dir} ${user}@${public_dns_name}:.${project_dir}
 echo "Copied project directory into instance"
 
 IFS=',' read -ra array <<< "${copy_paths}"
 for path in "${array[@]}"; do
 	# Make a directory and change the permissions
 	ssh -i ${pem_path} ${user}@${public_dns_name} "sudo mkdir -p .${path%/*}; sudo chmod 777 -R .${path%/*}"
+	exit_code=$?
+	if [ $exit_code -eq 1 ]; then
+		exit 1
+	fi
 
 	# Copy
 	scp -r -i ${pem_path} ${path} ${user}@${public_dns_name}:.${path%/*} 2> scp.log
@@ -82,17 +99,33 @@ for keyvalue in "${env_array[@]}"; do
 	# Update the key-value pair in .bashrc if it exists
     if ssh -i ${pem_path} ${user}@${public_dns_name} "grep -q '^export ${key}=' ~/.bashrc"; then
         ssh -i ${pem_path} ${user}@${public_dns_name} "sed -i 's/^export ${key}=.*$/export ${key}=${value}/' ~/.bashrc"
+		exit_code=$?
+		if [ $exit_code -eq 1 ]; then
+			exit 1
+		fi
 
     # Add the new key-value pair to the end of .bashrc if it doesn't exist
     else
         ssh -i ${pem_path} ${user}@${public_dns_name} "echo 'export ${key}=${value}' >> ~/.bashrc"
+		exit_code=$?
+		if [ $exit_code -eq 1 ]; then
+			exit 1
+		fi
     fi
 	echo "Updated environment variable ${key}=${value}"
 done
 
 # Reload .bashrc to update environment variables
 ssh -i ${pem_path} ${user}@${public_dns_name} "source ~/.bashrc"
+exit_code=$?
+if [ $exit_code -eq 1 ]; then
+	exit 1
+fi
 
 # Move all folders into the root folder
-ssh -i ${pem_path} ${user}@${public_dns_name} 'cd ~ && for dir in */; 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
+fi
 echo "Done updating remote project and file paths"

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):
         """

prism/tests/integration/test_run.py

@@ -1161,13 +1161,13 @@ def test_decorator_tasks_with_targets(self):
             "Sergey Prokopyev",
             "Dmitry Petelin",
             "Frank Rubio",
-            "Stephen Bowen",
-            "Warren Hoburg",
-            "Sultan Alneyadi",
-            "Andrey Fedyaev",
             "Jing Haiping",
             "Gui Haichow",
             "Zhu Yangzhu",
+            "Jasmin Moghbeli",
+            "Andreas Mogensen",
+            "Satoshi Furukawa",
+            "Konstantin Borisov",
         ]
         for n in names:
             formatted_name = n.lower().replace(" ", "_")