Git
/
ssh-audit
mirror of https://github.com/jtesta/ssh-audit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Wrap utils in single class.

This commit is contained in:
Andris Raugulis 2016-10-06 14:36:30 +03:00
parent 76f49d4016
commit 301a27ae27
1 changed files with 23 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
ssh-audit.py
View File

@ -133,7 +133,7 @@ class Output(object):
class OutputBuffer(list): class OutputBuffer(list):
def __enter__(self): def __enter__(self):
self.__buf = io.StringIO() self.__buf = utils.StringIO()
self.__stdout = sys.stdout self.__stdout = sys.stdout
sys.stdout = self.__buf sys.stdout = self.__buf
return self return self
@ -389,7 +389,7 @@ class SSH1(object):
class ReadBuf(object): class ReadBuf(object):
def __init__(self, data=None): def __init__(self, data=None):
super(ReadBuf, self).__init__() 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 self._len = len(data) if data else 0
@property @property
@ -1538,6 +1538,20 @@ def output(banner, header, kex=None, pkm=None):
output_fingerprint(kex, pkm, True, maxlen) output_fingerprint(kex, pkm, True, maxlen)
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): def parse_int(v):
try: try:
return int(v) return int(v)
@ -1576,7 +1590,7 @@ def parse_args():
s = args[0].split(':') s = args[0].split(':')
host, port = s[0].strip(), 22 host, port = s[0].strip(), 22
if len(s) > 1: if len(s) > 1:
port = parse_int(s[1]) port = utils.parse_int(s[1])
if not host or port <= 0: if not host or port <= 0:
usage('port {0} is not valid'.format(port)) usage('port {0} is not valid'.format(port))
conf.host = host conf.host = host
@ -1625,6 +1639,7 @@ def audit(conf, sshv=None):
output(banner, header, kex=kex) output(banner, header, kex=kex)
utils = Utils.wrap()
if __name__ == '__main__': if __name__ == '__main__':
out = Output() out = Output()
conf = parse_args() conf = parse_args()