#!/usr/bin/env python
"""Electrum Wallet Seed Generator (use at own risk).

Usage:

    -h --help
    -s --seed-bits=<bits>                   Size of the wallet seed in bits
                                                [default: 256]
    -i --iterations=<iterations>            Number of pbkdf2 iterations
                                                [default: 500000]
"""
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
import re
import sys
import hmac
import base64
import struct
import random
import hashlib
import operator as op
import itertools as it
import functools as ft
from collections import defaultdict

PY2 = sys.version_info.major == 2

if PY2:
    input = raw_input
    range = xrange
    map = it.imap
    zip = it.izip
    native_chr = chr
    chr = unichr

if PY2:
    ints2bytes = lambda ints: b"".join(map(native_chr, ints))
else:
    ints2bytes = bytes

DEFAULT_SEED_BITS = 256
PBKDF_ITERATIONS = 500000

# 256 words with large levenshtein distances from each other
WORDLIST = """
absolutely accidentally accounts acquaintance activities administration advertising
afghanistan afterwards against already alternative altogether ambassador ambulance
americans amsterdam anymore apologize apparently approaching appropriate architect
arrangement assassination atmosphere attitude autograph backyard barracks bartender
basketball bathroom beautifully because behaviour benefits boyfriend brazilian
breakfast businessman butterfly california casualties certificate charlie childhood
christian chuckling cigarettes circumstances cleveland collecting combination
comfortable commissioner community completed compromise concentrate congratulations
consciousness consequences considering conspiracy control coordinates countryside
creatures criminals curious dangerous defendant definitely democratic description
deserted developing difference difficulty disappear disappointed disgusting
distinguished disturb documentary downstairs earthquake electronic elevator
elizabeth elsewhere embarrassed emotional employees encourage enterprise enthusiasm
environment establish european everyone exactly exhibition expensive experienced
extraordinary farewell favorite fingerprints foreigners forgotten francisco
frankenstein frequency friendship frightened gentlemen goddamn graduation
grandchildren grandmother grandpa guaranteed hallelujah handcuffs handkerchief
handwriting happening headquarters helicopter helpful himself historical homework
humiliating husband identification ignorant immediately impossible incredibly
independence indicate indistinct individual inheritance instructor intellectual
intention interrupting introduce investigation involved journey kidnapping
kilometers knowledge leadership lemonade lifetime luck mademoiselle magnificent
maintenance manhattan meanwhile microphone millionaire minutes miserable
mississippi misunderstood motorcycle movement multiple naturally necessary
neighborhood nevertheless newspaper nicholas nothing occupied officer operations
opportunities orchestra orphanage overnight paintings particularly passion perfect
personally phenomenon philadelphia philosophy political portuguese possibility
practicing presentation previously problem professor pronounce psychiatrist
psychology punishment qualified questioning recommend refrigerator regardless
regular relationship remembering represents responsibilities restaurant revolution
ridiculous romantic sandwich satisfaction scholarship science screaming screenplay
sentimental shakespeare slippers sophisticated spaghetti special strawberry struggle
stubborn successful superintendent supermarket supervisor supposed surveillance
sweetheart switzerland sympathy technique telegram temperature territory thankyou
therapist thought tobacco"""


def levenshtein(string_1, string_2):
    if string_1 == string_2:
        return 0

    len_1 = len(string_1)
    len_2 = len(string_2)

    if len_1 == 0:
        return len_2
    if len_2 == 0:
        return len_1

    if len_1 > len_2:
        string_2, string_1 = string_1, string_2
        len_2, len_1 = len_1, len_2

    d0 = [i for i in range(len_2 + 1)]
    d1 = [j for j in range(len_2 + 1)]

    for i in range(len_1):
        d1[0] = i + 1
        for j in range(len_2):
            cost = d0[j]

            if string_1[i] != string_2[j]:
                cost += 1

                x_cost = d1[j] + 1
                if x_cost < cost:
                    cost = x_cost

                y_cost = d0[j + 1] + 1
                if y_cost < cost:
                    cost = y_cost

            d1[j + 1] = cost

        d0, d1 = d1, d0

    return d0[-1]


def pbkdf2(data, keylen, iterations=PBKDF_ITERATIONS, hashfunc=hashlib.sha256):
    mac = hmac.new(data, None, hashfunc)
    def _pseudorandom(x, mac=mac):
        h = mac.copy()
        h.update(x)
        if PY2:
            return map(ord, h.digest())
        return h.digest()

    buf = []
    for block in range(1, -(-keylen // mac.digest_size) + 1):
        rv = u = _pseudorandom(struct.pack(struct.pack(">I", block)))
        for i in range(iterations - 1):
            u = _pseudorandom(ints2bytes(u))
            rv = list(it.starmap(op.xor, zip(rv, u)))
        buf.extend(rv)

    return ints2bytes(buf[:keylen])


def dice_to_bytes(rolls):
    rolls = "".join([str(int(r) - 1) for r in rolls if r.isdigit()])
    hex_rolls = "{0:x}".format(int(rolls, 6))
    if len(hex_rolls) % 2 != 0:
        hex_rolls = "0" + hex_rolls
    hex_rolls = hex_rolls.upper().encode('ascii')
    return base64.b16decode(hex_rolls)


def normalized_words(text):
    return sorted(re.split("\s+", text.strip().lower()))[:256]


def clean_passphrase(phrase):
    words = normalized_words(WORDLIST)
    def lowest_clashes(word):
        clashes = defaultdict(list)
        for w in words:
            d = levenshtein(word, w)
            clashes[d].append(w)
        return clashes

    def guessword(word):
        clashes = lowest_clashes(word)
        if 0 in clashes:
            return clashes[0][0]
        for d in sorted(clashes.keys()):
            if len(clashes[d]) == 1:
                return clashes[d][0]
            return ""

    return [guessword(w) for w in normalized_words(phrase)]


def words_to_bytes(phrase_words):
    words = normalized_words(WORDLIST)
    word_indexes = (words.index(w) for w in phrase_words)
    return ints2bytes(word_indexes)


def random_phrase(length=4):
    words = normalized_words(WORDLIST)
    if PY2:
        digit_indexes = (ord(b) for b in os.urandom(length))
    else:
        digit_indexes = os.urandom(length)
    return " ".join((words[i] for i in digit_indexes))


def make_seed(
        dice_rolls,
        passphrase,
        seed_bits=DEFAULT_SEED_BITS,
        kdf=pbkdf2):
    target_bytes = seed_bits // 8

    dice_key = dice_to_bytes(dice_rolls)
    passphrase_bytes = words_to_bytes(clean_passphrase(passphrase))

    # ensure passphrase minimum entropy from passphrase
    dice_key = dice_key[:target_bytes - 8]
    keylen = target_bytes - len(dice_key)
    passphrase_key = kdf(passphrase_bytes, keylen=keylen)
    raw_seed = base64.b16encode(dice_key + passphrase_key)
    return raw_seed.decode('ascii').lower()


def getarg(argname, args):
    long_arg = '--' + argname
    short_arg = '-' + argname[:1]

    for idx, arg in enumerate(args):
        if arg.startswith(long_arg):
            if arg == long_arg:
                # check for parameter
                if idx + 1 < len(args):
                    nextarg = args[idx + 1]
                    if not nextarg.startswith("-"):
                        return nextarg
                return True
            if '=' in arg:
                argval = arg.split("=", 1)[-1]
                return argval
        if arg == short_arg:
            return True

    return


def test():
    dice_rolls = '1234561234' * 5
    words = normalized_words(WORDLIST)
    r = random.Random(0)
    passphrase = ' '.join([
        r.choice(words), r.choice(words), r.choice(words), r.choice(words)
    ])
    seed = make_seed(
        dice_rolls=dice_rolls,
        passphrase=passphrase,
        kdf=ft.partial(pbkdf2, iterations=8)
    )

    assert len(seed) == DEFAULT_SEED_BITS // 4
    assert seed == (
        "184ec43d79f1f70b8ce27e5e8043cdc32ac8eaff086cd55acb5cf6077fc4e1c7"
    )


def main(args=sys.argv[1:]):
    if getarg('help', args):
        print(__doc__.strip())
        return

    print(
        "Please roll a six sided dice 50x and write down your numbers "
        "on a peace of paper.\nIMPORTANT: You will need these numbers "
        "in order regenerate your wallet seed."
    )
    dice_rolls_1 = input("Dice Rolls: ")
    dice_rolls_2 = input("Repeat    : ")
    if dice_rolls_1 != dice_rolls_2:
        print("Missmatch of dice rolls")
        return

    passphrase = input("Passphrase (empty to generate one): ").strip()

    if not passphrase:
        passphrase = random_phrase()
        print(passphrase)

    seed_bits = int(getarg('seed-bits', args) or DEFAULT_SEED_BITS)
    assert seed_bits in (128, 256, 512)
    iterations = int(getarg('iterations', args) or PBKDF_ITERATIONS)

    seed = make_seed(
        dice_rolls=dice_rolls_1,
        passphrase=passphrase,
        seed_bits=seed_bits,
        kdf=ft.partial(pbkdf2, iterations=iterations)
    )
    print(seed)


if __name__ == '__main__':
    test()
    sys.exit(main() or 0)