From 301a27ae27c68562213efddf0f0fcc5f73c66e0a Mon Sep 17 00:00:00 2001 From: Andris Raugulis Date: Thu, 6 Oct 2016 14:36:30 +0300 Subject: [PATCH] Wrap utils in single class. --- ssh-audit.py | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/ssh-audit.py b/ssh-audit.py index e9243ec..0ae8111 100755 --- a/ssh-audit.py +++ b/ssh-audit.py @@ -133,7 +133,7 @@ class Output(object): class OutputBuffer(list): def __enter__(self): - self.__buf = io.StringIO() + self.__buf = utils.StringIO() self.__stdout = sys.stdout sys.stdout = self.__buf return self @@ -389,7 +389,7 @@ class SSH1(object): class ReadBuf(object): def __init__(self, data=None): super(ReadBuf, self).__init__() - self._buf = io.BytesIO(data) if data else io.BytesIO() + self._buf = utils.BytesIO(data) if data else utils.BytesIO() self._len = len(data) if data else 0 @property @@ -1538,11 +1538,25 @@ def output(banner, header, kex=None, pkm=None): output_fingerprint(kex, pkm, True, maxlen) -def parse_int(v): - try: - return int(v) - except: - return 0 +class Utils(object): + PY2 = sys.version_info[0] == 2 + + @classmethod + def wrap(cls): + o = cls() + if cls.PY2: + import StringIO + o.StringIO = o.BytesIO = StringIO.StringIO + else: + o.StringIO, o.BytesIO = io.StringIO, io.BytesIO + return o + + @staticmethod + def parse_int(v): + try: + return int(v) + except: + return 0 def parse_args(): @@ -1576,7 +1590,7 @@ def parse_args(): s = args[0].split(':') host, port = s[0].strip(), 22 if len(s) > 1: - port = parse_int(s[1]) + port = utils.parse_int(s[1]) if not host or port <= 0: usage('port {0} is not valid'.format(port)) conf.host = host @@ -1625,6 +1639,7 @@ def audit(conf, sshv=None): output(banner, header, kex=kex) +utils = Utils.wrap() if __name__ == '__main__': out = Output() conf = parse_args()