From 06992d7da60867d178d42a31b0de0808870398f1 Mon Sep 17 00:00:00 2001 From: Andris Raugulis Date: Fri, 1 Apr 2016 17:56:06 +0300 Subject: [PATCH] Refactor ssh connection within class for future improvements. --- ssh-audit.py | 133 +++++++++++++++++++++++++++++---------------------- 1 file changed, 76 insertions(+), 57 deletions(-) diff --git a/ssh-audit.py b/ssh-audit.py index f56f51c..d917d44 100755 --- a/ssh-audit.py +++ b/ssh-audit.py @@ -27,8 +27,6 @@ from __future__ import print_function import os, io, sys, socket, struct SSH_BANNER = 'SSH-2.0-OpenSSH_7.2' -SOCK_CONN_TIMEOUT = 3.0 -SOCK_READ_TIMEOUT = 5.0 def usage(): p = os.path.basename(sys.argv[0]) @@ -95,6 +93,7 @@ class Kex(object): class ReadBuf(object): def __init__(self, data = None): + super(ReadBuf, self).__init__() self._buf = io.BytesIO(data) if data else io.BytesIO() self._len = len(data) if data else 0 @@ -118,19 +117,69 @@ class ReadBuf(object): list_size = self.read_int() return self.read(list_size).decode().split(',') -class SockBuf(ReadBuf): - def __init__(self, s): - super(SockBuf, self).__init__() - self.__sock = s +class SSH(object): + MSG_KEXINIT = 20 + MSG_NEWKEYS = 21 + MSG_KEXDH_INIT = 30 + MSG_KEXDH_REPLY = 32 - def recv(self, size = 2048): - data = self.__sock.recv(size) - pos = self._buf.tell() - self._buf.seek(0, 2) - self._buf.write(data) - self._len += len(data) - self._buf.seek(pos, 0) - + class Socket(ReadBuf): + def __init__(self, host, port, cto = 3.0, rto = 5.0): + super(SSH.Socket, self).__init__() + try: + self.__sock = socket.create_connection((host, port), cto) + self.__sock.settimeout(rto) + except Exception as e: + out.fail('[fail] {}'.format(e)) + sys.exit(1) + + def __enter__(self): + return self + + def recv(self, size = 2048): + data = self.__sock.recv(size) + pos = self._buf.tell() + self._buf.seek(0, 2) + self._buf.write(data) + self._len += len(data) + self._buf.seek(pos, 0) + + def send(self, data): + self.__sock.send(data) + + def read_packet(self): + block_size = 8 + if self.unread_len < block_size: + self.recv() + header = self.read(block_size) + packet_size = struct.unpack('>I', header[:4])[0] + rest = header[4:] + lrest = len(rest) + padding = ord(rest[0:1]) + packet_type = ord(rest[1:2]) + if (packet_size - lrest) % block_size != 0: + out.fail('[exception] invalid ssh packet (block size)') + sys.exit(1) + rlen = packet_size - lrest + if self.unread_len < rlen: + self.recv() + buf = self.read(rlen) + packet = rest[2:] + buf[0:packet_size - lrest] + payload = packet[0:packet_size - padding] + return packet_type, payload + + def __del__(self): + self.__cleanup() + + def __exit__(self, ex_type, ex_value, tb): + self.__cleanup() + + def __cleanup(self): + try: + self.__sock.shutdown(socket.SHUT_RDWR) + self.__sock.close() + except: + pass def get_ssh_ver(versions): tv = [] @@ -309,26 +358,6 @@ def process_kex(kex): process_algorithms('mac', kex.server.mac, maxlen) out.sep() -def read_ssh_packet(sbuf): - block_size = 8 - if sbuf.unread_len < block_size: - sbuf.recv() - header = sbuf.read(block_size) - packet_size = struct.unpack('>I', header[:4])[0] - rest = header[4:] - lrest = len(rest) - padding = ord(rest[0:1]) - packet_type = ord(rest[1:2]) - if (packet_size - lrest) % block_size != 0: - out.fail('[exception] invalid ssh packet (block size)') - sys.exit(1) - rlen = packet_size - lrest - if sbuf.unread_len < rlen: - sbuf.recv() - buf = sbuf.read(rlen) - packet = rest[2:] + buf[0:packet_size - lrest] - payload = packet[0:packet_size - padding] - return packet_type, payload def parse_int(v): try: @@ -355,30 +384,20 @@ def parse_args(): def main(): host, port = parse_args() - s = None - try: - s = socket.create_connection((host, port), SOCK_CONN_TIMEOUT) - s.settimeout(SOCK_READ_TIMEOUT) - sbuf = SockBuf(s) - s.send(SSH_BANNER.encode() + b'\r\n') - sbuf.recv() - banner = sbuf.read_line() - out.head('# general') - out.good('[info] banner: ' + banner) - if banner.startswith('SSH-1.99-'): - out.fail('[fail] protocol SSH1 enabled') - packet_type, payload = read_ssh_packet(sbuf) - if packet_type != 20: - out.fail('[exception] did not receive MSG_KEXINIT (20), instead received unknown message ({0})'.format(packet_type)) - sys.exit(1) - kex = Kex.parse(payload) - process_kex(kex) - except Exception as e: - out.fail('[fail] {}'.format(e)) + s = SSH.Socket(host, port) + s.send(SSH_BANNER.encode() + b'\r\n') + s.recv() + banner = s.read_line() + out.head('# general') + out.good('[info] banner: ' + banner) + if banner.startswith('SSH-1.99-'): + out.fail('[fail] protocol SSH1 enabled') + packet_type, payload = s.read_packet() + if packet_type != SSH.MSG_KEXINIT: + out.fail('[exception] did not receive MSG_KEXINIT (20), instead received unknown message ({0})'.format(packet_type)) sys.exit(1) - finally: - if s: - s.close() + kex = Kex.parse(payload) + process_kex(kex) if __name__ == '__main__': out = Output()