advent-of-code-day4.py

import md5


def find_key(secret):
    """
    >>> find_key('abcdef')
    609043
    >>> find_key('pqrstuv')
    1048970
    """
    key = ''
    current_number = 1
    while key[0:5] != '00000':
        key = md5.new(secret + str(current_number)).hexdigest()
        current_number += 1
    return current_number - 1


def find_key_p2(secret):
    key = ''
    current_number = 1
    while key[0:6] != '000000':
        key = md5.new(secret + str(current_number)).hexdigest()
        current_number += 1
    return current_number - 1


if __name__ == "__main__":
    import doctest
    doctest.testmod()
    print 'Answer is: ' + str(find_key('yzbqklnj'))
    print 'Second Answer is: ' + str(find_key_p2('yzbqklnj'))