Add python2 compatibility, use defusedxml

remip2 · remip2 · commit 3c26112722e4 · 2019-02-25T17:16:48.000+01:00
diff --git a/.gitignore b/.gitignore
@@ -4,3 +4,4 @@ build
 .eggs
 stormshield.sns.sslclient.egg-info
 __pycache__
+*.pyc
diff --git a/README.md b/README.md
@@ -116,7 +116,7 @@ Concerning the SSL validation:
 
 ## Tests
 
-Warning: tests require a remote SNS appliance.
+Warning: some tests require a remote SNS appliance.
 
 `$ PASSWORD=password APPLIANCE=10.0.0.254 python3 setup.py test`
 
diff --git a/bin/snscli b/bin/snscli
@@ -1,15 +1,16 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 
 """ cli to connect to Stormshield Network Security appliances"""
 
+from __future__ import unicode_literals
 import sys
 import os
 import re
 import logging
 import readline
 import getpass
 import atexit
-import xml.dom.minidom
+import defusedxml.minidom
 import begin
 from pygments import highlight
 from pygments.lexers import XmlLexer
@@ -47,20 +48,18 @@ def make_completer():
         return results[state]
     return custom_complete
 
-@begin.start(auto_convert=True, short_args=False, lexical_order=True)
-@begin.logging
-def main(host: 'Remote UTM' = None,
-         ip: 'Remote UTM ip' = None,
-         usercert: 'User certificate file' = None,
-         cabundle: 'CA bundle file' = None,
-         password: 'Password' = None,
-         port: 'Remote port' = 443,
-         user: 'User name' = 'admin',
-         sslverifypeer: 'Strict SSL CA check' = True,
-         sslverifyhost: 'Strict SSL host name check' = True,
-         credentials: 'Privilege list' = None,
-         script: 'Command script' = None,
-         outputformat: 'Output format (ini|xml)' = 'ini'):
+def main(host = None,
+         ip = None,
+         usercert = None,
+         cabundle = None,
+         password = None,
+         port = 443,
+         user = 'admin',
+         sslverifypeer = True,
+         sslverifyhost = True,
+         credentials = None,
+         script = None,
+         outputformat = 'ini'):
 
     for handler in logging.getLogger().handlers:
         if handler.__class__ == logging.StreamHandler:
@@ -123,7 +122,7 @@ def main(host: 'Remote UTM' = None,
                 logging.error(str(exception))
                 sys.exit(1)
             if outputformat == 'xml':
-                print(highlight(xml.dom.minidom.parseString(response.xml).toprettyxml(),
+                print(highlight(defusedxml.minidom.parseString(response.xml).toprettyxml(),
                                 XmlLexer(), TerminalFormatter()))
             else:
                 print(response.output)
@@ -191,7 +190,30 @@ def main(host: 'Remote UTM' = None,
                 logging.error(str(exception))
         else:
             if outputformat == 'xml':
-                print(highlight(xml.dom.minidom.parseString(response.xml).toprettyxml(),
+                print(highlight(defusedxml.minidom.parseString(response.xml).toprettyxml(),
                                 XmlLexer(), TerminalFormatter()))
             else:
                 print(response.output)
+
+# set function parameter annotations manually for compatibility with python2
+main.__annotations__ = {
+            'host': 'Remote UTM',
+            'ip': 'Remote UTM ip',
+            'usercert': 'User certificate file',
+            'cabundle': 'CA bundle file',
+            'password': 'Password',
+            'port': 'Remote port',
+            'user': 'User name',
+            'sslverifypeer': 'Strict SSL CA check',
+            'sslverifyhost': 'Strict SSL host name check',
+            'credentials': 'Privilege list',
+            'script': 'Command script',
+            'outputformat': 'Output format (ini|xml)'
+        }
+# use correct input function with python2
+try:
+    input = raw_input
+except NameError:
+    pass
+
+begin.start(begin.logging(main))
diff --git a/setup.py b/setup.py
@@ -26,12 +26,15 @@
         'requests',
         'requests_toolbelt',
         'colorlog',
-        'pyreadline; platform_system == "Windows"'
+        'defusedxml',
+        'pyreadline; platform_system == "Windows"',
+        'py2-ipaddress; python_version < "3"'
     ],
     include_package_data=True,
     tests_require=["nose"],
     test_suite='nose.collector',
     classifiers=[
+        "Programming Language :: Python :: 2",
         "Programming Language :: Python :: 3",
         "License :: Apache License 2.0",
         "Operating System :: OS Independent",
diff --git a/stormshield/sns/__init__.py b/stormshield/sns/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "1.0.0.beta"
+__version__ = "1.0.0.beta1"
 
 # major.minor.patch
 # major: breaking API change
diff --git a/stormshield/sns/configparser.py b/stormshield/sns/configparser.py
@@ -7,6 +7,7 @@
 in ini/section format.
 """
 
+import sys
 import re
 from shlex import shlex
 from requests.structures import CaseInsensitiveDict
@@ -50,10 +51,10 @@ def __init__(self, text):
         lines = text.splitlines()
 
         # strip serverd headers if needed
-        group = self.SERVERD_HEAD_RE.match(lines[0])
-        if group:
+        match = self.SERVERD_HEAD_RE.match(lines[0])
+        if match:
             del lines[0]
-            self.format = group[1]
+            self.format = match.group(1)
         if self.SERVERD_TAIL_RE.match(lines[-1]):
             del lines[-1]
 
@@ -76,9 +77,9 @@ def __init__(self, text):
                 continue
 
             # section header
-            group = self.SECTION_RE.match(line)
-            if group:
-                section = group[1]
+            match = self.SECTION_RE.match(line)
+            if match:
+                section = match.group(1)
                 if self.format == 'section':
                     self.data[section] = CaseInsensitiveDict()
                 else:
@@ -88,19 +89,22 @@ def __init__(self, text):
             if self.format == "list":
                 self.data[section].append(line)
             elif self.format == "section_line":
+                # fix encoding for python2
+                if sys.version_info[0] < 3:
+                    line = line.encode('utf-8')
                 # parse token=value token2=value2
                 lexer = shlex(line, posix=True)
                 lexer.wordchars += "=.-*:,"
                 parsed = {}
                 for word in lexer:
                     # ignore anything else than token=value
                     if '=' in word:
-                        token, value = word.split("=", maxsplit=1)
+                        token, value = word.split("=", 1)
                         parsed[token] = value
                 self.data[section].append(parsed)
             else:
                 # section
-                (token, value) = line.split("=", maxsplit=1)
+                (token, value) = line.split("=", 1)
                 self.data[section][token] = unquote(value)
 
 
diff --git a/stormshield/sns/crc.py b/stormshield/sns/crc.py
@@ -82,6 +82,10 @@
 def compute_crc32(data):
     """ Returns the hex string of CRC value of DATA """
 
+    # python2 convert str to bytes
+    if type(data) == str:
+        data = bytearray(data)
+
     n = len(data)
     crc = CRC32_init
 
@@ -93,6 +97,10 @@ def compute_crc32(data):
 def update_crc32(data, crc):
     """ Incremental CRC, returns updated CRC. """
 
+    # python2 convert str to bytes
+    if type(data) == str:
+        data = bytearray(data)
+
     n = len(data)
 
     for i in range(0, n):
diff --git a/stormshield/sns/sslclient.py b/stormshield/sns/sslclient.py
@@ -1,4 +1,5 @@
 #!/usr/bin/python
+# -*- coding: utf-8 -*-
 
 """
 stormshield.sns.sslclient
@@ -8,13 +9,14 @@
 and Response class to handle API answers.
 """
 
+from __future__ import unicode_literals
 import os
 import ipaddress
 import base64
 import logging
 import re
 import platform
-import xml.etree.ElementTree as ElementTree
+import defusedxml.ElementTree as ElementTree
 import requests
 from requests.adapters import HTTPAdapter, DEFAULT_POOLSIZE, DEFAULT_RETRIES, DEFAULT_POOLBLOCK
 from requests.packages.urllib3.poolmanager import PoolManager
@@ -339,7 +341,7 @@ def connect(self):
         self.logger.log(logging.DEBUG, request.text)
 
         try:
-            nws_node = ElementTree.fromstring(request.text)
+            nws_node = ElementTree.fromstring(request.content)
             msg = nws_node.attrib['msg']
         except:
             raise AuthenticationError("Can't decode authentication result")
@@ -359,7 +361,7 @@ def connect(self):
         self.logger.log(logging.DEBUG, request.text)
 
         if request.status_code == requests.codes.OK:
-            nws_node = ElementTree.fromstring(request.text)
+            nws_node = ElementTree.fromstring(request.content)
             ret = int(nws_node.attrib['code'])
             msg = nws_node.attrib['msg']
 
@@ -421,13 +423,13 @@ def send_command(self, command):
 
         request = self.session.get(
             self.baseurl + '/api/command?sessionid=' + self.sessionid +
-            '&cmd=' + requests.compat.quote(command), # manually done since we need %20 encoding
+            '&cmd=' + requests.compat.quote(command.encode('utf-8')), # manually done since we need %20 encoding
             headers=self.headers)
 
         self.logger.log(logging.DEBUG, request.text)
 
         if request.status_code == requests.codes.OK:
-            nws_node = ElementTree.fromstring(request.text)
+            nws_node = ElementTree.fromstring(request.content)
             code = int(nws_node.attrib['code'])
             self.nws_parse(code)
             serverd = nws_node[0]
@@ -440,7 +442,7 @@ def send_command(self, command):
                 response = Response(ret=serverd_ret,
                                     code=serverd_code,
                                     msg=serverd_msg,
-                                    output=format_output(request.text),
+                                    output=format_output(request.content),
                                     xml=request.text)
 
                 #multiline answer get the final code
@@ -523,14 +525,14 @@ def upload(self, filename):
         uploadh.close()
 
         if request.status_code == requests.codes.OK:
-            nws_node = ElementTree.fromstring(request.text)
+            nws_node = ElementTree.fromstring(request.content)
             code = int(nws_node.attrib['code'])
             self.nws_parse(code)
 
             return Response(code=nws_node[0].get('code'),
                             ret=int(nws_node[0].get('ret')),
                             msg=nws_node[0].get('msg'),
-                            output=format_output(request.text),
+                            output=format_output(request.content),
                             xml=request.text)
 
         raise ServerError("HTTP error {}".format(request.status_code))
diff --git a/tests/test_auth.py b/tests/test_auth.py
@@ -4,7 +4,7 @@
 import unittest
 from stormshield.sns.sslclient import SSLClient
 
-APPLIANCE=os.getenv('APPLIANCE', "")
+APPLIANCE = os.getenv('APPLIANCE', "")
 SERIAL = os.getenv('SERIAL', "")
 PASSWORD = os.getenv('PASSWORD', "")
 
diff --git a/tests/test_cert.py b/tests/test_cert.py
@@ -4,7 +4,7 @@
 import unittest
 from stormshield.sns.sslclient import SSLClient
 
-APPLIANCE=os.getenv('APPLIANCE', "")
+APPLIANCE = os.getenv('APPLIANCE', "")
 FQDN = os.getenv('FQDN', "")
 PASSWORD = os.getenv('PASSWORD', "")
 CABUNDLE = os.getenv('CABUNDLE', "")
@@ -26,11 +26,11 @@ def test_sslverifypeer(self):
             client = SSLClient(host=APPLIANCE, user='admin', password=PASSWORD)
             self.fail("SSLClient should have failed (untrusted CA)")
         except Exception as exception:
-            self.assertTrue(1==1, "SSLClient did not connect (untrusted CA)")
+            self.assertTrue(True, "SSLClient did not connect (untrusted CA)")
         
         try:
             client = SSLClient(host=APPLIANCE, user='admin', password=PASSWORD, sslverifypeer=False)
-            self.assertTrue(1==1, "SSLClient connects with sslverifypeer=True")
+            self.assertTrue(True, "SSLClient connects with sslverifypeer=False")
         except Exception as exception:
             print(exception)
             self.fail("SSLClient did not connect")
diff --git a/tests/test_configparser.py b/tests/test_configparser.py
@@ -1,4 +1,5 @@
 #!/usr/bin/python
+# -*- coding: utf-8 -*-
 
 import unittest
 import json
diff --git a/tests/test_file.py b/tests/test_file.py
@@ -1,14 +1,17 @@
 #!/usr/bin/python
+# -*- coding: utf-8 -*-
 
+from __future__ import unicode_literals
 import random
 import string
 import os
 import sys
 import tempfile
 import unittest
+import shutil
 from stormshield.sns.sslclient import SSLClient
 
-APPLIANCE=os.getenv('APPLIANCE', "")
+APPLIANCE = os.getenv('APPLIANCE', "")
 PASSWORD = os.getenv('PASSWORD', "")
 
 @unittest.skipIf(APPLIANCE=="", "APPLIANCE env var must be set to the ip/hostname of a running SNS appliance")
@@ -19,20 +22,22 @@ class TestFormatIni(unittest.TestCase):
     def setUp(self):
         self.client = SSLClient(host=APPLIANCE, user='admin', password=PASSWORD, sslverifyhost=False)
 
-        self.tmpdir = tempfile.TemporaryDirectory()
-        self.upload = os.path.join(self.tmpdir.name, 'upload')
-        self.download = os.path.join(self.tmpdir.name, 'download')
+        self.tmpdir = tempfile.mkdtemp()
+        self.upload = os.path.join(self.tmpdir, 'upload')
+        self.download = os.path.join(self.tmpdir, 'download')
 
     def tearDown(self):
         self.client.disconnect()
-        self.tmpdir.cleanup()
+        shutil.rmtree(self.tmpdir, ignore_errors=True)
 
     def test_upload_download(self):
         """ Test file upload and download """
 
+        letters = string.ascii_letters + 'éèàÎîô'
+
         #generate a random file
-        content = "".join( [random.choice(string.ascii_letters) for i in range(15)] )
-        with open(self.upload, "w") as fh:
+        content = ("".join( [random.choice(letters) for i in range(2500)] )).encode('utf-8')
+        with open(self.upload, "wb") as fh:
             fh.write(content)
 
         response = self.client.send_command('CONFIG COMMUNICATION EMAIL TEMPLATE UPLOAD pvm_summary < ' + self.upload)
@@ -43,7 +48,7 @@ def test_upload_download(self):
         response = self.client.send_command('CONFIG COMMUNICATION EMAIL TEMPLATE DOWNLOAD pvm_summary > ' + self.download)
         self.assertEqual(response.ret, 100)
 
-        with open(self.download, "r") as fh:
+        with open(self.download, "rb") as fh:
             downloaded = fh.read()
 
         self.assertEqual(content, downloaded)
diff --git a/tests/test_format_ini.py b/tests/test_format_ini.py
@@ -6,7 +6,7 @@
 
 from stormshield.sns.sslclient import SSLClient
 
-APPLIANCE=os.getenv('APPLIANCE', "")
+APPLIANCE = os.getenv('APPLIANCE', "")
 PASSWORD = os.getenv('PASSWORD', "")
 
 @unittest.skipIf(APPLIANCE=="", "APPLIANCE env var must be set to the ip/hostname of a running SNS appliance")
diff --git a/tests/test_utf8.py b/tests/test_utf8.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-__version__ = "1.0.0.beta"`
	`1`	`+__version__ = "1.0.0.beta1"`
`2`	`2`
`3`	`3`	`# major.minor.patch`
`4`	`4`	`# major: breaking API change`