#! /usr/bin/env python

# Convert OpenSSH known_hosts and known_hosts2 files to "new format" PuTTY
# host keys.
#   usage:
#     kh2reg.py [ --win ] known_hosts1 2 3 4 ... > hosts.reg
#       Creates a Windows .REG file (double-click to install).
#     kh2reg.py --unix    known_hosts1 2 3 4 ... > sshhostkeys
#       Creates data suitable for storing in ~/.putty/sshhostkeys (Unix).
# Line endings are someone else's problem as is traditional.
# Originally developed for Python 1.5.2, but probably won't run on that
# any more.

import fileinput
import base64
import struct
import string
import re
import sys
import getopt
import itertools
import collections

def winmungestr(s):
    "Duplicate of PuTTY's mungestr() in winstore.c:1.10 for Registry keys"
    candot = 0
    r = ""
    for c in s:
        if c in ' \*?%~' or ord(c)<ord(' ') or (c == '.' and not candot):
            r = r + ("%%%02X" % ord(c))
        else:
            r = r + c
        candot = 1
    return r

def strtolong(s):
    "Convert arbitrary-length big-endian binary data to a Python long"
    bytes = struct.unpack(">%luB" % len(s), s)
    return reduce ((lambda a, b: (long(a) << 8) + long(b)), bytes)

def strtolong_le(s):
    "Convert arbitrary-length little-endian binary data to a Python long"
    bytes = reversed(struct.unpack(">%luB" % len(s), s))
    return reduce ((lambda a, b: (long(a) << 8) + long(b)), bytes)

def longtohex(n):
    """Convert long int to lower-case hex.

    Ick, Python (at least in 1.5.2) doesn't appear to have a way to
    turn a long int into an unadorned hex string -- % gets upset if the
    number is too big, and raw hex() uses uppercase (sometimes), and
    adds unwanted "0x...L" around it."""

    plain=string.lower(re.match(r"0x([0-9A-Fa-f]*)l?$", hex(n), re.I).group(1))
    return "0x" + plain

def warn(s):
    "Warning with file/line number"
    sys.stderr.write("%s:%d: %s\n"
                     % (fileinput.filename(), fileinput.filelineno(), s))

output_type = 'windows'

def invert(n, p):
    """Compute inverse mod p."""
    if n % p == 0:
        raise ZeroDivisionError()
    a = n, 1, 0
    b = p, 0, 1
    while b[0]:
        q = a[0] // b[0]
        a = a[0] - q*b[0], a[1] - q*b[1], a[2] - q*b[2]
        b, a = a, b
    assert abs(a[0]) == 1
    return a[1]*a[0]

def jacobi(n,m):
    """Compute the Jacobi symbol.

    The special case of this when m is prime is the Legendre symbol,
    which is 0 if n is congruent to 0 mod m; 1 if n is congruent to a
    non-zero square number mod m; -1 if n is not congruent to any
    square mod m.

    """
    assert m & 1
    acc = 1
    while True:
        n %= m
        if n == 0:
            return 0
        while not (n & 1):
            n >>= 1
            if (m & 7) not in {1,7}:
                acc *= -1
        if n == 1:
            return acc
        if (n & 3) == 3 and (m & 3) == 3:
            acc *= -1
        n, m = m, n

class SqrtModP(object):
    """Class for finding square roots of numbers mod p.

    p must be an odd prime (but its primality is not checked)."""

    def __init__(self, p):
        p = abs(p)
        assert p & 1
        self.p = p

        # Decompose p as 2^e k + 1 for odd k.
        self.k = p-1
        self.e = 0
        while not (self.k & 1):
            self.k >>= 1
            self.e += 1

        # Find a non-square mod p.
        for self.z in itertools.count(1):
            if jacobi(self.z, self.p) == -1:
                break
        self.zinv = invert(self.z, self.p)

    def sqrt_recurse(self, a):
        ak = pow(a, self.k, self.p)
        for i in range(self.e, -1, -1):
            if ak == 1:
                break
            ak = ak*ak % self.p
        assert i > 0
        if i == self.e:
            return pow(a, (self.k+1) // 2, self.p)
        r_prime = self.sqrt_recurse(a * pow(self.z, 2**i, self.p))
        return r_prime * pow(self.zinv, 2**(i-1), self.p) % self.p

    def sqrt(self, a):
        j = jacobi(a, self.p)
        if j == 0:
            return 0
        if j < 0:
            raise ValueError("{} has no square root mod {}".format(a, self.p))
        a %= self.p
        r = self.sqrt_recurse(a)
        assert r*r % self.p == a
        # Normalise to the smaller (or 'positive') one of the two roots.
        return min(r, self.p - r)

    def __str__(self):
        return "{}({})".format(type(self).__name__, self.p)
    def __repr__(self):
        return self.__str__()

    instances = {}

    @classmethod
    def make(cls, p):
        if p not in cls.instances:
            cls.instances[p] = cls(p)
        return cls.instances[p]

    @classmethod
    def root(cls, n, p):
        return cls.make(p).sqrt(n)

NistCurve = collections.namedtuple("NistCurve", "p a b")
nist_curves = {
    "ecdsa-sha2-nistp256": NistCurve(0xffffffff00000001000000000000000000000000ffffffffffffffffffffffff, 0xffffffff00000001000000000000000000000000fffffffffffffffffffffffc, 0x5ac635d8aa3a93e7b3ebbd55769886bc651d06b0cc53b0f63bce3c3e27d2604b),
    "ecdsa-sha2-nistp384": NistCurve(0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffeffffffff0000000000000000ffffffff, 0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffeffffffff0000000000000000fffffffc, 0xb3312fa7e23ee7e4988e056be3f82d19181d9c6efe8141120314088f5013875ac656398d8a2ed19d2a85c8edd3ec2aef),
    "ecdsa-sha2-nistp521": NistCurve(0x01ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff, 0x01fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffc, 0x0051953eb9618e1c9a1f929a21a0b68540eea2da725b99b315f3b8b489918ef109e156193951ec7e937b1652c0bd3bb1bf073573df883d2c34f1ef451fd46b503f00),
}

try:
    optlist, args = getopt.getopt(sys.argv[1:], '', [ 'win', 'unix' ])
    if filter(lambda x: x[0] == '--unix', optlist):
        output_type = 'unix'
except getopt.error, e:
    sys.stderr.write(str(e) + "\n")
    sys.exit(1)

if output_type == 'windows':
    # Output REG file header.
    sys.stdout.write("""REGEDIT4

[HKEY_CURRENT_USER\Software\SimonTatham\PuTTY\SshHostKeys]
""")

class BlankInputLine(Exception):
    pass

class UnknownKeyType(Exception):
    def __init__(self, keytype):
        self.keytype = keytype

class KeyFormatError(Exception):
    def __init__(self, msg):
        self.msg = msg

# Now process all known_hosts input.
for line in fileinput.input(args):

    try:
        # Remove leading/trailing whitespace (should zap CR and LF)
        line = string.strip (line)

        # Skip blanks and comments
        if line == '' or line[0] == '#':
            raise BlankInputLine

        # Split line on spaces.
        fields = string.split (line, ' ')

        # Common fields
        hostpat = fields[0]
        keyparams = []      # placeholder
        keytype = ""        # placeholder

        # Grotty heuristic to distinguish known_hosts from known_hosts2:
        # is second field entirely decimal digits?
        if re.match (r"\d*$", fields[1]):

            # Treat as SSH-1-type host key.
            # Format: hostpat bits10 exp10 mod10 comment...
            # (PuTTY doesn't store the number of bits.)
            keyparams = map (long, fields[2:4])
            keytype = "rsa"

        else:

            # Treat as SSH-2-type host key.
            # Format: hostpat keytype keyblob64 comment...
            sshkeytype, blob = fields[1], base64.decodestring (fields[2])

            # 'blob' consists of a number of
            #   uint32    N (big-endian)
            #   uint8[N]  field_data
            subfields = []
            while blob:
                sizefmt = ">L"
                (size,) = struct.unpack (sizefmt, blob[0:4])
                size = int(size)   # req'd for slicage
                (data,) = struct.unpack (">%lus" % size, blob[4:size+4])
                subfields.append(data)
                blob = blob [struct.calcsize(sizefmt) + size : ]

            # The first field is keytype again.
            if subfields[0] != sshkeytype:
                raise KeyFormatError("""
                    outer and embedded key types do not match: '%s', '%s'
                    """ % (sshkeytype, subfields[1]))

            # Translate key type string into something PuTTY can use, and
            # munge the rest of the data.
            if sshkeytype == "ssh-rsa":
                keytype = "rsa2"
                # The rest of the subfields we can treat as an opaque list
                # of bignums (same numbers and order as stored by PuTTY).
                keyparams = map (strtolong, subfields[1:])

            elif sshkeytype == "ssh-dss":
                keytype = "dss"
                # Same again.
                keyparams = map (strtolong, subfields[1:])

            elif sshkeytype in nist_curves:
                keytype = sshkeytype
                # Have to parse this a bit.
                if len(subfields) > 3:
                    raise KeyFormatError("too many subfields in blob")
                (curvename, Q) = subfields[1:]
                # First is yet another copy of the key name.
                if not re.match("ecdsa-sha2-" + re.escape(curvename),
                                sshkeytype):
                    raise KeyFormatError("key type mismatch ('%s' vs '%s')"
                            % (sshkeytype, curvename))
                # Second contains key material X and Y (hopefully).
                # First a magic octet indicating point compression.
                point_type = struct.unpack("B", Q[0])[0]
                Qrest = Q[1:]
                if point_type == 4:
                    # Then two equal-length bignums (X and Y).
                    bnlen = len(Qrest)
                    if (bnlen % 1) != 0:
                        raise KeyFormatError("odd-length X+Y")
                    bnlen = bnlen // 2
                    x = strtolong(Qrest[:bnlen])
                    y = strtolong(Qrest[bnlen:])
                elif 2 <= point_type <= 3:
                    # A compressed point just specifies X, and leaves
                    # Y implicit except for parity, so we have to
                    # recover it from the curve equation.
                    curve = nist_curves[sshkeytype]
                    x = strtolong(Qrest)
                    yy = (x*x*x + curve.a*x + curve.b) % curve.p
                    y = SqrtModP.root(yy, curve.p)
                    if y % 2 != point_type % 2:
                        y = curve.p - y

                keyparams = [curvename, x, y]

            elif sshkeytype == "ssh-ed25519":
                keytype = sshkeytype

                if len(subfields) != 2:
                    raise KeyFormatError("wrong number of subfields in blob")
                if subfields[0] != sshkeytype:
                    raise KeyFormatError("key type mismatch ('%s' vs '%s')"
                            % (sshkeytype, subfields[0]))
                # Key material y, with the top bit being repurposed as
                # the expected parity of the associated x (point
                # compression).
                y = strtolong_le(subfields[1])
                x_parity = y >> 255
                y &= ~(1 << 255)

                # Standard Ed25519 parameters.
                p = 2**255 - 19
                d = 0x52036cee2b6ffe738cc740797779e89800700a4d4141d8ab75eb4dca135978a3

                # Recover x^2 = (y^2 - 1) / (d y^2 + 1).
                xx = (y*y - 1) * invert(d*y*y + 1, p) % p

                # Take the square root.
                x = SqrtModP.root(xx, p)

                # Pick the square root of the correct parity.
                if (x % 2) != x_parity:
                    x = p - x

                keyparams = [x, y]
            else:
                raise UnknownKeyType(sshkeytype)

        # Now print out one line per host pattern, discarding wildcards.
        for host in string.split (hostpat, ','):
            if re.search (r"[*?!]", host):
                warn("skipping wildcard host pattern '%s'" % host)
                continue
            elif re.match (r"\|", host):
                warn("skipping hashed hostname '%s'" % host)
                continue
            else:
                m = re.match (r"\[([^]]*)\]:(\d*)$", host)
                if m:
                    (host, port) = m.group(1,2)
                    port = int(port)
                else:
                    port = 22
                # Slightly bizarre output key format: 'type@port:hostname'
                # XXX: does PuTTY do anything useful with literal IP[v4]s?
                key = keytype + ("@%d:%s" % (port, host))
                # Most of these are numbers, but there's the occasional
                # string that needs passing through
                value = string.join (map (
                    lambda x: x if isinstance(x, basestring) else longtohex(x),
                    keyparams), ',')
                if output_type == 'unix':
                    # Unix format.
                    sys.stdout.write('%s %s\n' % (key, value))
                else:
                    # Windows format.
                    # XXX: worry about double quotes?
                    sys.stdout.write("\"%s\"=\"%s\"\n"
                                     % (winmungestr(key), value))

    except UnknownKeyType, k:
        warn("unknown SSH key type '%s', skipping" % k.keytype)
    except KeyFormatError, k:
        warn("trouble parsing key (%s), skipping" % k.msg)
    except BlankInputLine:
        pass

# The spec at http://support.microsoft.com/kb/310516 says we need
# a blank line at the end of the reg file:
#
#   Note the registry file should contain a blank line at the
#   bottom of the file.
#
if output_type == 'windows':
    # Output REG file header.
    sys.stdout.write("\n")