Support for specifying client IP address, or automatically getting it…

… from ipify (#114)

* add ip_address and auto_determine_ip_address config options

* add auto_determine_ip_address to example config in readme

* add docstring and type hint to get_ip_address

* add ipify to pipfile

* add ipify to setup.py

* update tests

* add test for get_ip_address

Pipfile

@@ -9,6 +9,7 @@ verify_ssl = true
 "boto3" = "*"
 onelogin = "*"
 keyring = "*"
+ipify = "*"
 
 [requires]
 python_version = ">=3.5"

README.md

@@ -162,6 +162,12 @@ and can define any additional directives as desired.
 - `otp_device` - Allow the automatic selection of an OTP device.  
   This value is the human readable string name for the device.
   Eg, `OneLogin Protect`, `Yubico YubiKey`, etc
+- `ip_address` - The client IP address to send to OneLogin.
+  Relevant when using OneLogin Policies with an IP whitelist.
+  If this is specified, `auto_determine_ip_address` is not used.
+- `auto_determine_ip_address` - Automatically determine the client IP address.
+  Relevant when using OneLogin Policies with an IP whitelist.
+  Can be used without specifying `ip_address`.
 
 ### Example
 
@@ -175,6 +181,7 @@ client_secret = a85234b6db01a29a493e2422d7930dffe6f4d3a826270a18838574f6b8ef7c3e
 save_password = yes
 profile = mycompany-onelogin
 duration_seconds = 3600
+auto_determine_ip_address = yes
 
 [testing]
 aws_app_id = 555029

onelogin_aws_cli/__init__.py

@@ -1,11 +1,16 @@
 """OneLogin/AWS Business logic"""
 
+from typing import Optional
+
 import configparser
 import xml.etree.ElementTree as ElementTree
 
 import base64
 import boto3
 import os
+
+import ipify
+
 from onelogin.api.client import OneLoginClient
 
 from onelogin_aws_cli.configuration import Section
@@ -49,17 +54,19 @@ def get_saml_assertion(self):
         self.user_credentials.load_credentials()
 
         saml_resp = self.ol_client.get_saml_assertion(
-            self.user_credentials.username,
-            self.user_credentials.password,
-            self.config['aws_app_id'],
-            self.config['subdomain']
+            username_or_email=self.user_credentials.username,
+            password=self.user_credentials.password,
+            app_id=self.config['aws_app_id'],
+            subdomain=self.config['subdomain'],
+            ip_address=self.get_ip_address(),
         )
 
         if saml_resp is None:
             raise Exception("Onelogin Error: '{error}' '{desc}'".format(
                 error=self.ol_client.error,
                 desc=self.ol_client.error_description
             ))
+
         if saml_resp.mfa:
             if not self.mfa.ready():
                 self.mfa.select_device(saml_resp.mfa.devices)
@@ -75,6 +82,24 @@ def get_saml_assertion(self):
 
         self.saml = saml_resp
 
+    def get_ip_address(self) -> Optional[str]:
+        """
+        Get the client IP address.
+        Uses either the `ip_address` in config,
+        or if `auto_determine_ip_address` is specified in config,
+        the ipify service is used to dynamically lookup the IP address.
+        """
+
+        # if ip address has been hard coded in config file, use that
+        ip_address = self.config.get('ip_address')
+        if ip_address is not None:
+            return ip_address
+
+        # if auto determine is enabled, use ipify to lookup the ip
+        if self.config.auto_determine_ip_address:
+            ip_address = ipify.get_ip()
+            return ip_address
+
     def get_arns(self):
         """Extract the IAM Role ARNs from the SAML Assertion"""
 

onelogin_aws_cli/configuration.py

@@ -127,6 +127,13 @@ def can_save_password(self) -> bool:
             fallback=self.config.DEFAULTS['save_password']
         )
 
+    @property
+    def auto_determine_ip_address(self):
+        return self.config.getboolean(
+            self.section_name,
+            "auto_determine_ip_address",
+        )
+
     def set_overrides(self, overrides: dict):
         """
         Specify a dictionary values which take precedence over the existing

onelogin_aws_cli/tests/test_oneloginAWS.py

@@ -27,7 +27,9 @@ def setUp(self):
             client_id='mock-id',
             client_secret='mock-secret',
             username='mock-username',
-            duration_seconds=2600
+            duration_seconds=2600,
+            ip_address='1.2.3.4',
+            auto_determine_ip_address=False
         ))
         self.ol_with_role = OneloginAWS(dict(
             base_uri="https://api.us.onelogin.com/",
@@ -51,6 +53,12 @@ def test_init(self):
         self.assertEqual(mock_config, ol.config)
         self.assertEqual('mock-username', ol.user_credentials.username)
 
+    def test_get_ip_address(self):
+        self.ol.saml = Namespace(saml_response=self.SAML_SINGLE_ROLE)
+        ip_address = self.ol.get_ip_address()
+
+        self.assertEqual('1.2.3.4', ip_address)
+
     def test_get_arns(self):
         self.ol.saml = Namespace(saml_response=self.SAML_SINGLE_ROLE)
         self.ol.get_arns()

onelogin_aws_cli/tests/test_oneloginAWS_saml.py

@@ -33,7 +33,8 @@ def setUp(self):
                 subdomain='example',
                 can_save_password=False,
                 username='mock-username',
-                duration_seconds=2600
+                duration_seconds=2600,
+                auto_determine_ip_address=False,
             ),
         )
 
@@ -62,8 +63,11 @@ def test_get_saml_assertion_single(self, getpw):
         self.assertEqual(self.ol.saml, 'mock-saml-response')
 
         self.get_saml_assertion_mock.assert_called_with(
-            'mock-username', 'mock-password',
-            'mock-app-id', 'example'
+            username_or_email='mock-username',
+            password='mock-password',
+            app_id='mock-app-id',
+            subdomain='example',
+            ip_address=None,
         )
 
         self.get_saml_assertion_verifying_mock.assert_called_with(
@@ -93,8 +97,11 @@ def test_get_saml_assertion_multiple(self, getpw):
         self.assertEqual(self.ol.saml, 'mock-saml-response')
 
         self.get_saml_assertion_mock.assert_called_with(
-            'mock-username', 'mock-password',
-            'mock-app-id', 'example'
+            username_or_email='mock-username',
+            password='mock-password',
+            app_id='mock-app-id',
+            subdomain='example',
+            ip_address=None,
         )
 
         self.get_saml_assertion_verifying_mock.assert_called_with(
@@ -121,7 +128,8 @@ def test_username_unspecified(self):
                 aws_app_id='mock-app-id',
                 subdomain='example',
                 can_save_password=False,
-                duration_seconds=2600
+                duration_seconds=2600,
+                auto_determine_ip_address=False,
             ),
         )
         self.assertIsNone(ol.user_credentials.username)

requirements.txt

@@ -1,3 +1,4 @@
 boto3
 onelogin
 keyring
+ipify

setup.py

@@ -32,7 +32,8 @@
     install_requires=[
         'boto3',
         'onelogin',
-        'keyring'
+        'keyring',
+        'ipify',
     ],
     setup_requires=['nose>=1.0'],
     entry_points={