diff --git a/docs/source/docs/providers/aws/instance-parameters.md b/docs/source/docs/providers/aws/instance-parameters.md index 7dd93d2..e8d5a1a 100644 --- a/docs/source/docs/providers/aws/instance-parameters.md +++ b/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. diff --git a/setup.py b/setup.py index a633101..4bfc484 100644 --- a/setup.py +++ b/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', diff --git a/spotty/providers/aws/cfn_templates/instance/data/template.yaml b/spotty/providers/aws/cfn_templates/instance/data/template.yaml index b999b4b..2e492d0 100644 --- a/spotty/providers/aws/cfn_templates/instance/data/template.yaml +++ b/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 diff --git a/spotty/providers/aws/cfn_templates/instance/template.py b/spotty/providers/aws/cfn_templates/instance/template.py index f46afc2..ddf9094 100644 --- a/spotty/providers/aws/cfn_templates/instance/template.py +++ b/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 diff --git a/spotty/providers/aws/config/instance_config.py b/spotty/providers/aws/config/instance_config.py index fca01fc..e21dee4 100644 --- a/spotty/providers/aws/config/instance_config.py +++ b/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'] diff --git a/spotty/providers/aws/config/validation.py b/spotty/providers/aws/config/validation.py index 5b3eb2b..beaeb8b 100644 --- a/spotty/providers/aws/config/validation.py +++ b/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 = [ diff --git a/tests/providers/aws/config/instance_config_validation.py b/tests/providers/aws/config/instance_config_validation.py index 29e1d64..0d7aec2 100644 --- a/tests/providers/aws/config/instance_config_validation.py +++ b/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))