AWS: Inbound IP whitelist

docs/source/docs/providers/aws/instance-parameters.md

@@ -65,6 +65,9 @@ Read more: [Caching Docker Image on an EBS Volume].
     ```
     It will open ports 6006 for TensorBoard and 8888 for Jupyter Notebook. 
 
+-__`inboundIp`__ _(optional)_ - an IP address or CIDR range to use as a whitelist for all ports provided
+to the security group. The default is to not use a whitelist. 
+
 - __`localSshPort`__ _(optional)_ - if this parameter is set, all the Spotty commands will create SSH connections 
 with the instance using the IP address __127.0.0.1__ and the specified port. This can be useful in case when an 
 instance doesn't have a public IP address and a jump-server is used for tunneling.

setup.py

@@ -49,6 +49,7 @@ def get_description():
           'cfn_flip',  # to work with CloudFormation templates
           'schema',
           'chevron',
+          'netaddr' # validates configuration IP addresses
       ],
       tests_require=['moto'],
       test_suite='tests',

spotty/providers/aws/cfn_templates/instance/data/template.yaml

@@ -88,15 +88,7 @@ Resources:
           IpProtocol: -1
           FromPort: 0
           ToPort: 65535
-      SecurityGroupIngress:
-        - CidrIp: 0.0.0.0/0
-          IpProtocol: tcp
-          FromPort: 22
-          ToPort: 22
-        - CidrIpv6: ::/0
-          IpProtocol: tcp
-          FromPort: 22
-          ToPort: 22
+      SecurityGroupIngress: []
 
   PreparingInstanceSignal:
     Type: AWS::CloudFormation::WaitCondition

spotty/providers/aws/cfn_templates/instance/template.py

@@ -56,20 +56,16 @@ def prepare_instance_template(ec2, instance_config: InstanceConfig, docker_comma
             }]
         del template['Resources']['InstanceLaunchTemplate']['Properties']['LaunchTemplateData']['SecurityGroupIds']
 
-    # add ports to the security group
-    for port in instance_config.ports:
-        if port != 22:
-            template['Resources']['InstanceSecurityGroup']['Properties']['SecurityGroupIngress'] += [{
-                'CidrIp': '0.0.0.0/0',
-                'IpProtocol': 'tcp',
-                'FromPort': port,
-                'ToPort': port,
-            }, {
-                'CidrIpv6': '::/0',
-                'IpProtocol': 'tcp',
-                'FromPort': port,
-                'ToPort': port,
-            }]
+    # add ports to the security group for the user-defined inbound IP
+    inbound_ip = instance_config.inbound_ip
+    cidr_ip = 'CidrIpv6' if inbound_ip.version == 6 else 'CidrIp'
+    for port in [22] + instance_config.ports:
+        template['Resources']['InstanceSecurityGroup']['Properties']['SecurityGroupIngress'] += [{
+            cidr_ip: str(inbound_ip),
+            'IpProtocol': 'tcp',
+            'FromPort': port,
+            'ToPort': port,
+        }]
 
     if instance_config.is_spot_instance:
         # set maximum price

spotty/providers/aws/config/instance_config.py

@@ -75,7 +75,11 @@ def max_price(self) -> float:
     @property
     def managed_policy_arns(self) -> list:
         return self._params['managedPolicyArns']
-    
+
     @property
     def instance_profile_arn(self) -> str:
         return self._params['instanceProfileArn']
+
+    @property
+    def inbound_ip(self) -> str:
+        return self._params['inboundIp']

spotty/providers/aws/config/validation.py

@@ -1,5 +1,6 @@
 import os
 from schema import Schema, Optional, And, Regex, Or, Use
+from netaddr import IPNetwork
 from spotty.config.validation import validate_config, get_instance_parameters_schema, has_prefix
 
 
@@ -30,6 +31,8 @@ def validate_instance_parameters(params: dict):
                                              ),
         Optional('managedPolicyArns', default=[]): [str],
         Optional('instanceProfileArn', default=None): str,
+        Optional('inboundIp', default='0.0.0.0/0'): Use(IPNetwork, error='"inboundIp" should be a valid IP '
+                                                                         'address or CIDR range'),
     }
 
     volumes_checks = [

tests/providers/aws/config/instance_config_validation.py

@@ -28,6 +28,7 @@ def test_default_configuration(self):
             'spotInstance': False,
             'subnetId': '',
             'volumes': [],
+            'inboundIp': '0.0.0.0/0',
         }
 
         self.assertEqual(expected_params, validate_instance_parameters(required_params))