Merge pull request #15 from deep-security/integration-v2.1

Integration v2.1.1 of deep-security-py

Author: marknca

lib/core.py

@@ -10,14 +10,12 @@
 # project libraries
 import lib.deepsecurity as deepsecurity
 
-def get_arg_parser(prog='ds-to-aws-waf.py', description=None, add_help=False):
+def get_arg_parser(prog='ds-cli.py', description=None, add_help=False):
   """
   Create a standardized argument parser
   """
   if not description:
-    description = """
-    Create and update AWS WAF WACL rules based on information from a Deep Security installation
-"""
+    description = "A command line interface to Trend Micro's Deep Security"
 
   parser = argparse.ArgumentParser(prog=prog, description=description, add_help=add_help)
 
@@ -33,6 +31,11 @@ def get_arg_parser(prog='ds-to-aws-waf.py', description=None, add_help=False):
   parser.add_argument('-s', '--aws-secret-key', action='store', dest='aws_secret_key', required=False, help='The secret key for an IAM identity in the AWS account to connect to')
   parser.add_argument('-r', '--aws-region', action='store', dest='aws_region', required=False, default='us-east-1', help='The name of AWS region to connect to')
 
+  # Azure arguments
+  parser.add_argument('--azure-subscription-id', action='store', dest='azure_subscription_id', required=False, help='The subscription ID of the Azure subscription to connect to')
+  parser.add_argument('--azure-certificate', action='store', dest='azure_certificate', required=False, help='The certificate associated with the Azure subscription to connect to')
+  parser.add_argument('--azure-certificate-password', action='store', dest='azure_certificate_password', required=False, default='us-east-1', help='The password associated with the specified Azure certificate')
+  
   # general structure arguments
   parser.add_argument('--ignore-ssl-validation', action='store_true', dest='ignore_ssl_validation', required=False, help='Ignore SSL certification validation. Be careful when you use this as it disables a recommended security check. Required for Deep Security Managers using a self-signed SSL certificate')
   parser.add_argument('--dryrun', action='store_true', required=False, help='Do a dry run of the command. This will not make any changes to your AWS WAF service')

lib/deepsecurity/__init__.py

@@ -6,4 +6,6 @@
 sys.path.append(current_path)
 
 # import project files as required
-import dsm
+import dsm
+import translation
+translation.Terms.read_terms_file()

lib/deepsecurity/core.py

@@ -76,11 +76,12 @@ def _set_logging(self):
 
     return logger
 
-  def _get_request_format(self, api=None, call=None):
+  def _get_request_format(self, api=None, call=None, use_cookie_auth=False):
     if not api: api = self.API_TYPE_SOAP
     return {
       'api': api,
       'call': call,
+      'use_cookie_auth': use_cookie_auth,
       'query': None,
       'data': None,
     }
@@ -108,6 +109,9 @@ def _request(self, request, auth_required=True):
         REST API calls this will be a dict converted to JSON automatically 
         by this method
 
+      use_cookie_auth
+        Whether or not to use an HTTP Cookie in lieu of a querystring for authorization
+
     ## Output
 
     Returns a dict:
@@ -139,9 +143,9 @@ def _request(self, request, auth_required=True):
     # add the authentication parameters
     if auth_required:
       if request['api'] == self.API_TYPE_REST:
-        # sID is a query string
-        if not request['query']: request['query'] = {}
-        request['query']['sID'] = self._sessions[self.API_TYPE_REST]
+        if not request['use_cookie_auth']: # sID is a query string
+          if not request['query']: request['query'] = {}
+          request['query']['sID'] = self._sessions[self.API_TYPE_REST]
       elif request['api'] == self.API_TYPE_SOAP:
         # sID is part of the data
         if not request['data']: request['data'] = {}
@@ -182,6 +186,11 @@ def _request(self, request, auth_required=True):
 
     # authentication calls don't accept the Accept header
     if request['call'].startswith('authentication'): del(headers['Accept'])
+    
+    # some rest calls use a cookie to pass the sID
+    if request['api'] == self.API_TYPE_REST and request['use_cookie_auth']:
+      headers['Cookie'] = 'sID="{}"'.format(self._sessions[self.API_TYPE_REST])
+
     if request['api'] == self.API_TYPE_REST and request['call'] in [
       'apiVersion',
       'status/manager/ping'
@@ -229,6 +238,7 @@ def _request(self, request, auth_required=True):
     result = {
       'status': response.getcode() if response else None,
       'raw': response.read() if response else None,
+      'headers': dict(response.headers) if response else dict(),
       'data': None
     }
     bytes_of_data = len(result['raw']) if result['raw'] else 0
@@ -254,12 +264,14 @@ def _request(self, request, auth_required=True):
       else:
         # JSON response
         try:
-          if result['raw']:
-            result['data'] = json.loads(result['raw'])
+          if result['raw'] and result['status'] != 204:
+            result['type'] = result['headers']['content-type']
+            result['data'] = json.loads(result['raw']) if 'json' in result['type'] else None
         except Exception, json_err:
           # report the exception as 'info' because it's not fatal and the data is 
           # still captured in result['raw']
           self.log("Could not convert response from call {} to JSON. Threw exception:\n\t{}".format(request['call'], json_err), level='info')
+          
     return result
 
   def _prefix_keys(self, prefix, d):
@@ -270,7 +282,7 @@ def _prefix_keys(self, prefix, d):
     if not type(d) == type({}): return d
     new_d = d.copy()
     for k,v in d.items():
-      new_key = "{}:{}".format(prefix, k)
+      new_key = u"{}:{}".format(prefix, k)
       new_v = v
       if type(v) == type({}): new_v = self._prefix_keys(prefix, v)
       new_d[new_key] = new_v 
@@ -283,7 +295,7 @@ def _prep_data_for_soap(self, call, details):
     Prepare the complete XML SOAP envelope
     """
     data = xmltodict.unparse(self._prefix_keys('ns1', { call: details }), pretty=False, full_document=False)
-    soap_xml = """
+    soap_xml = u"""
     <?xml version="1.0" encoding="UTF-8"?>
     <SOAP-ENV:Envelope xmlns:ns0="http://schemas.xmlsoap.org/soap/envelope/" xmlns:ns1="urn:Manager" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xmlns:SOAP-ENV="http://schemas.xmlsoap.org/soap/envelope/">
       <SOAP-ENV:Header/>
@@ -293,7 +305,10 @@ def _prep_data_for_soap(self, call, details):
     </SOAP-ENV:Envelope>
     """.format(data).strip()
 
-    return soap_xml
+    # convert any nil values to the proper format
+    soap_xml = re.sub(r'<([^>]+)></\1>', r'<\1 xsi:nil="true" />', soap_xml)
+
+    return soap_xml.encode('utf-8')
 
   def log(self, message='', err=None, level='info'):
     """
@@ -432,4 +447,112 @@ def _set_properties(self, api_response, log_func):
         setattr(self, new_key, val)
       except Exception, err:
         if log_func:
-          log_func("Could not set property {} to value {} for object {}".format(k, v, s))
+          log_func("Could not set property {} to value {} for object {}".format(k, v, s))
+          try:
+            setattr(self, log, log_func)
+          except: pass
+
+  def to_dict(self):
+    """
+    Convert the object properties to API keypairs
+    """
+    result = {}
+
+    api_properties = translation.Terms.api_to_new.values()
+
+    for p in dir(self):
+      if p in api_properties:
+        key = translation.Terms.get_reverse(p)
+        val = getattr(self, p)
+        result[key] = val
+
+    return result
+
+class CoreList(list):
+  def __init__(self, *args):
+    super(CoreList, self).__init__(args)
+    self._exempt_from_find = []
+
+  def find(self, **kwargs):
+    """
+    Find any items where the values match the cumulative kwargs patterns and return their indices
+
+    If a keyword's value is a list, .find will match on any value for that keyword
+
+    .find(id=1)
+    >>> returns any item with a property 'id' and value in [1]
+        possibilities:
+           { 'id': 1, 'name': 'One'}
+           { 'id': 1, 'name': 'Two'}
+        
+    .find(id=[1,2])
+    >>> returns any item with a property 'id' and value in [1,2]
+        possibilities:
+           { 'id': 1, 'name': 'One'}
+           { 'id': 2, 'name': 'One'}
+           { 'id': 1, 'name': 'Two'}
+           { 'id': 2, 'name': 'Two'}
+
+    .find(id=1, name='One')
+    >>> returns any item with a property 'id' and value in [1] AND a property 'name' and value in ['One']
+        possibilities:
+           { 'id': 1, 'name': 'One'}
+        
+    .find(id=[1,2], name='One')
+    >>> returns any item with a property 'id' and value in [1,2] AND a property 'name' and value in ['One']
+        possibilities:
+           { 'id': 1, 'name': 'One'}
+           { 'id': 2, 'name': 'One'}
+
+    .find(id=[1,2], name=['One,Two'])
+    >>> returns any item with a property 'id' and value in [1,2] AND a property 'name' and value in ['One','Two']
+        possibilities:
+           { 'id': 1, 'name': 'One'}
+           { 'id': 2, 'name': 'One'}
+           { 'id': 1, 'name': 'Two'}
+           { 'id': 2, 'name': 'Two'}
+    """
+    results = []
+
+    if kwargs:
+      for item_id, item in enumerate(self):
+        item_matches = False
+        for match_attr, match_attr_vals in kwargs.items():
+          if not type(match_attr_vals) == type([]): match_attr_vals = [match_attr_vals]
+
+          # does the current item have the property
+          attr_to_check = None
+          if match_attr in dir(item):
+            attr_to_check = getattr(item, match_attr)
+          elif 'has_key' in dir(item) and item.has_key(match_attr):
+            attr_to_check = item[match_attr]
+
+          if attr_to_check:
+            # does the property match the specified values?
+            for match_attr_val in match_attr_vals:
+              if type(attr_to_check) in [type(''), type(u'')]:
+                # string comparison
+                match = re.search(r'{}'.format(match_attr_val), attr_to_check)
+                if match:
+                  item_matches = True
+                  break # and move on to the new kwarg
+                else:
+                  item_matches = False
+              elif type(attr_to_check) == type([]):
+                # check for the match in the list
+                if match_attr_val in attr_to_check:
+                  item_matches = True
+                  break # and move on to the new kwarg
+                else:
+                  item_matches = False
+              else:
+                # object comparison
+                if attr_to_check == match_attr_val:
+                  item_matches = True
+                  break # and move on to the new kwarg
+                else:
+                  item_matches = False
+
+        if item_matches: results.append(item_id)
+
+    return results

lib/deepsecurity/dsm.py

@@ -18,14 +18,16 @@ def __init__(self,
       tenant=None,
       username=None,
       password=None,
+      prefix="",
       ignore_ssl_validation=False
       ):
     core.CoreApi.__init__(self)
     self._hostname = None
     self._port = port
-    self._tenant = tenant
-    self._username = username
-    self._password = password
+    self._tenant = unicode(tenant, "utf-8") if tenant else None # no harm in converting to ensure compatibility with non-latin tenant names
+    self._username = unicode(username, "utf-8") if username else None
+    self._password = unicode(password, "utf-8") if password else None
+    self._prefix = prefix
     self.ignore_ssl_validation = ignore_ssl_validation
     self.hostname = hostname
 
@@ -48,7 +50,8 @@ def __str__(self):
     """
     Return a better string representation
     """
-    return "Manager <{}:{}>".format(self.hostname, self.port)
+    dsm_port = ":{}".format(self.port) if self.port else ""
+    return "Manager <{}{}>".format(self.hostname, dsm_port)
 
   # *******************************************************************
   # properties
@@ -68,7 +71,7 @@ def port(self): return self._port
 
   @port.setter
   def port(self, value):
-    self._port = int(value)
+    self._port = int(value) if value else None
     self._set_endpoints()
 
   @property
@@ -95,15 +98,24 @@ def password(self, value):
     self._password = value
     self._reset_session()
 
+  @property
+  def prefix(self): return self._prefix
+
+  @prefix.setter
+  def prefix(self, value):
+    if not value or not type(value) in [type(''), type(u'')]: value = ""
+    self._prefix = value
+
   # *******************************************************************
   # methods
   # *******************************************************************
   def _set_endpoints(self):
     """
     Set the API endpoints based on the current configuration
     """
-    self._rest_api_endpoint = "https://{}:{}/rest".format(self.hostname, self.port)
-    self._soap_api_endpoint = "https://{}:{}/webservice/Manager".format(self.hostname, self.port)
+    dsm_port = ":{}".format(self.port) if self.port else "" # allow for endpoints with no port specified
+    self._rest_api_endpoint = "https://{}{}/{}rest".format(self.hostname, dsm_port, self.prefix)
+    self._soap_api_endpoint = "https://{}{}/{}webservice/Manager".format(self.hostname, dsm_port, self.prefix)
 
   def _reset_session(self):
     """
@@ -170,10 +182,10 @@ def sign_out(self):
       response = self._request(rest_call)
       if response and response['status'] == 200: self._sessions[self.API_TYPE_REST] = None
 
-    if self._sessions[self.API_TYPE_REST] and self._sessions[self.API_TYPE_SOAP]:
-      return True
-    else:
+    if self._sessions[self.API_TYPE_REST] or self._sessions[self.API_TYPE_SOAP]:
       return False
+    else:
+      return True
 
   def get_api_version(self):
     """

lib/deepsecurity/environments.py

@@ -31,20 +31,25 @@ def add_aws_account(self, name, aws_access_key=None, aws_secret_key=None, region
     responses = {}
 
     regions = {
-      'us-east-1': 'amazon.cloud.region.key.1',
-      'us-west-1': 'amazon.cloud.region.key.2',
-      'us-west-2': 'amazon.cloud.region.key.3',
-      'eu-west-1': 'amazon.cloud.region.key.4',
-      'ap-southeast-1': 'amazon.cloud.region.key.5',
-      'ap-northeast-1': 'amazon.cloud.region.key.6',
-      'sa-east-1': 'amazon.cloud.region.key.7',
+      'us-east-1': 'amazon.cloud.region.key.1', # N. Virginia
+      'us-west-1': 'amazon.cloud.region.key.2', # N. California
+      'us-west-2': 'amazon.cloud.region.key.3', # Oregon
+      'eu-west-1': 'amazon.cloud.region.key.4', # Ireland
+      'ap-southeast-1': 'amazon.cloud.region.key.5', # Singapore
+      'ap-northeast-1': 'amazon.cloud.region.key.6', # Tokyo
+      'sa-east-1': 'amazon.cloud.region.key.7', # Sao Paulo
+      'ap-southeast-2': 'amazon.cloud.region.key.8', # Sydney
+      # need to add:
+      #   ap-south-1 / Mumbai
+      #   ap-northeast-2 / Seoul
+      #   eu-central-1 / Frankfurt
     }
 
     regions_to_add = []
     if regions.has_key(region):
       regions_to_add.append(region)
     elif region == 'all':
-      regions_to_add.append(regions.keys())
+      regions_to_add = regions.keys()
     else:
       self.log("A region must be specified when add an AWS account to Deep Security")