#!/usr/bin/python

import math
import string
import sys

def get_transformations():
    '''
    Return two dictionaries for substring substitution: reverse, forward
    
    `reverse` will map passwordish input to possible lowercase outputs.
    `forward` will map lowercase input to possible passwordish outputs.
    
    Ex:
        reverse['$'] -> ['s']
        forward['s'] -> ['S', 's', 'Z', 'z', '$']
    '''
    class listdict(dict):
        def __missing__(self, key):
            self[key] = []
            return self[key]
    # Dictionary of substring subsition 
    transformations = listdict({
        "ate": ["8",],
        "for": ["4",],
        "to": ["2",],
        "too": ["2",],
        "a": ["4", "@", "/-\\", "^",],
        "c": ["k",],
        "e": ["3",],
        "i": ["1", "l", "|",],
        "l": ["1", "l", "|",],
        "o": ["0",],
        "s": ["$", "z", "Z", "5",],
        "t": ["7",],
        "z": ["s", "S", "2",],
    })
    for key in range(ord("a"), ord("z")+1):
        transformations[chr(key)].append(chr(key))
        transformations[chr(key)].append(chr(key).upper())
    reverse_lookup = listdict()
    for key in transformations:
        for value in transformations[key]:
            reverse_lookup[value].append(key)
    return dict(reverse_lookup), dict(transformations)


def get_word_list():
    '''
    Return a list of lowercase English words that are long enough.
    '''
    s = open('/usr/share/dict/words').read() + 'troubador'
    words = set()
    for line in s.split('\n'):
        if len(line) > 3 and "'" not in line:
            words.add(line.lower())
    return words


def measure_word(password, comb=1, word_list=None):
    '''
    Decide if `password` begins with an English word.
    
    Return: None
        Password does not begin with an English word.
        
    Return: bits, remaining_password
        Password does begin with an English word.
        `bits` is the entropy of the measuered part of the password.
        `remaining_password` is the unmeasured part (after the word).
    '''
    BITS_PER_CHAR = 1.9
    BONUS_BITS = 1.0
    
    if word_list is None:
        word_list = get_word_list()
    
    candidate = None
    for key in sorted(reverse, key=lambda x: len(x), reverse=True):
        if password.startswith(key):
            for value in reverse[key]:
                new_word_list = []
                for word in word_list:
                    if word == value:
                        # Must set to variable
                        # Ex.  "hello" would otherwise match "hell"
                        # because "l" is in the word list, but so is "lo" which does
                        # not match "l".
                        # Set the variable but keep searching. In the next level,
                        # "o" is in the word list and that's the time to return.
                        candidate = math.log(comb * 2**BONUS_BITS, 2), password[len(key):]
                    if word.startswith(value):
                        new_word_list.append(word[len(value):])
                if new_word_list:
                    if key == value:
                        comb_factor = 1
                    else:
                        comb_factor = len(forward[value])
                    buf = measure_word(
                        password[len(key):],
                        comb * comb_factor * 2**BITS_PER_CHAR,
                        new_word_list
                    )
                    if buf is not None:
                        return buf
    else:
        # No better match has been found.
        return candidate


def measure(password):
    '''
    Return how many bits of entropy there are in the password.
    '''
    MAGIC_FACTOR = 0.6
    bits = 0
    misc_chars = ''
    words_cache = {}
    chars_cache = set()
    while password:
        # Check if it is known to not be a word.
        for char in chars_cache:
            if password[0] == char:
                password = password[1:]
                misc_chars += char
                break
        else:
            # Check if it is known to be a word.
            for word in words_cache:
                if password.startswith(word):
                    password = password[len(word):]
                    words_cache[word] += 1
                    break
            else:
                # Not cached, measure word.
                tmp = measure_word(password)
                if tmp is not None:
                    wordlen = len(password) - len(tmp[1])
                    word = password[:wordlen]
                    words_cache[word] = 1
                    # First instance of word.
                    bits += tmp[0]
                    password = tmp[1]
                else:
                    # Not a word.
                    misc_chars += password[0]
                    chars_cache.add(password[0])
                    password = password[1:]
    # Repeated words.
    for key in words_cache:
        bits += math.log(words_cache[key], 2)
    # Deal with misc_chars.
    if misc_chars:
        types = {
            "upper": string.uppercase,
            "lower": string.lowercase,
            "digits": string.digits,
            "punctuation": string.punctuation,
        }
        # 26 + 26 + 10 + something
        expected_set = 0
        for key in types:
            value = filter(lambda x: x in types[key], misc_chars)
            if value:
                expected_set += len(types[key])
        if expected_set == 0:
            expected_set = 128
        # Decide whether the miscellaneous characters are random or not.
        poor = False
        if len(misc_chars) > expected_set:
            for ch in set(misc_chars):
                if float(misc_chars.count(ch))/len(misc_chars) > 2.0/expected_set:
                    poor = True
                    break
        else:
            if float(len(misc_chars))/len(set(misc_chars)) > 2.0:
                poor = True
        if poor:
            bits += (
                # Repetitive
                math.log(len(misc_chars))
                # Predictable set
                * math.log(len(set(misc_chars)), 2)
                # xkcd factor for "Tr0ub4dor&3"
                * MAGIC_FACTOR
                # Some credit for at least trying
                + 4
            )
        else:
            bits += len(misc_chars) * math.log(expected_set, 2) * MAGIC_FACTOR
    return bits


def main():
    messages = {
        10: "Are you kidding me?",
        20: "Vulnerable",
        30: "Can't you do better than that?",
        40: "Almost good enough",
        50: "Good enough",
        60: "Good",
        70: "Very good",
        80: "Secure",
        90: "Awesome",
    }
    try:
        while True:
            sys.stderr.write('\nMeasure password strength: ')
            sys.stderr.flush()
            
            line = sys.stdin.readline()
            if not line:
                break
            bits = measure(line.rstrip('\n'))
            
            sys.stdout.write('~{} bits: '.format(int(bits + 0.5)))
            for level in sorted(messages):
                if level > bits:
                    sys.stdout.write(messages[level] + '\n')
                    break
            else:
                sys.stdout.write('Speachless\n\n')
    except KeyboardInterrupt:
        pass
    sys.stderr.write('\n')



reverse, forward = get_transformations()

if __name__ == '__main__':
    main()