#!/usr/bin/env python3
import argparse
from nullroute.core import Core
from nullroute.io import SshPacketReader
from pprint import pprint
import io
import socket
import sys

# Binary packet protocol:
#   https://tools.ietf.org/html/rfc4253#section-6
# KEXINIT packet format:
#   https://tools.ietf.org/html/rfc4253#section-7.1

SSH_MSG_DISCONNECT = 1
SSH_MSG_IGNORE = 2
SSH_MSG_DEBUG = 4
SSH_MSG_KEXINIT = 20

def read_ssh_packets(r):
    br = SshPacketReader(r)
    while True:
        Core.trace("reading binary packet...")
        packet_len = br.read_uint32()
        padding_len = br.read_byte()
        payload = br.read(packet_len - padding_len - 1)
        padding = br.read(padding_len)
        mac = br.read(0)
        # no MAC for initial handshake packets

        Core.trace("parsed binary packet:")
        Core.trace("  packet_len = %d", packet_len)
        Core.trace("  padding_len = %d", padding_len)
        Core.trace("  payload = [%d] %r", len(payload), payload)
        Core.trace("  padding = [%d] %r", len(padding), padding)
        Core.trace("  mac = [%d] %r", len(mac), mac)
        yield payload, mac

def show_kex_packet(br):
    def _parse_kex_packet(br):
        u_fields = ["Key exchange", "Host key signature"]
        b_fields = ["Cipher", "MAC", "Compression", "Languages"]
        cookie = br.read(16)
        for field in u_fields:
            values = br.read_array()
            yield field, values
        for field in b_fields:
            values_c2s = br.read_array()
            values_s2c = br.read_array()
            if values_c2s == values_s2c:
                yield field, values_c2s
            else:
                yield f"{field} (c2s)", values_c2s
                yield f"{field} (s2c)", values_s2c

    for field, values in _parse_kex_packet(br):
        print(f"{field}:")
        for v in values:
            print(f"    \"{v.decode()}\"")
        if not values:
            print(f"    (empty list)")

def dump_one(buf, name):
    print("\033[1m%s:\033[m" % name)
    words = buf.decode().split(",")
    pprint(words)

def dump_hdr(buf):
    packet = buf
    while packet:
        banner, packet = packet.split(b"\n", 1)
        if banner.startswith(b"SSH-"):
            break
    Core.debug("parsed buffer:")
    Core.debug("  banner = [%d] %r", len(banner), banner)
    Core.debug("  packet = [%d] %r", len(packet), packet)
    print("SSH version banner: %r" % banner)

    for payload, *_ in read_ssh_packets(io.BytesIO(packet)):
        br = SshPacketReader(io.BytesIO(payload))
        packet_type = br.read_byte()
        if packet_type == SSH_MSG_DISCONNECT:
            Core.debug("got SSH_MSG_DISCONNECT packet")
            reason = br.read_uint32()
            text = br.read_string().decode("utf-8", errors="ignore")
            lang = br.read_string()
            raise IOError("received a disconnect packet: reason=%d, %r" % (reason, text))
        elif packet_type == SSH_MSG_IGNORE:
            Core.debug("got SSH_MSG_IGNORE packet")
            data = br.read_string()
            Core.debug("  data = [%d] %r", len(data), data)
        elif packet_type == SSH_MSG_DEBUG:
            Core.debug("got SSH_MSG_DEBUG")
            display = br.read_bool()
            text = br.read_string().decode("utf-8", errors="ignore")
            lang = br.read_string()
            Core.debug("  text = [%d] %r", len(text), text)
        elif packet_type == SSH_MSG_KEXINIT:
            Core.debug("got SSH_MSG_KEXINIT")
            show_kex_packet(br)
            break
        else:
            raise IOError("received message was not SSH_MSG_KEXINIT (got %d)" % packet_type)

parser = argparse.ArgumentParser()
parser.add_argument("-p", "--port", type=int, default=22)
parser.add_argument("host")
args = parser.parse_args()

gais = socket.getaddrinfo(args.host, args.port, type=socket.SOCK_STREAM)

for (g_family, g_type, g_proto, g_cname, g_addr) in gais:
    sock = socket.socket(g_family, g_type, g_proto)
    sock.connect(g_addr)
    sock.send(b"SSH-2.0-Cipherscan\r\n")
    # Dropbear
    buf = sock.recv(64)
    buf += sock.recv(16*1024)
    sock.close()
    dump_hdr(buf)
    break
