# Author: Andrieiev Danil Maksimovich
# danssg08@gmail.com | https://github.com/DanilAndreev

from collections import Callable
from random import Random
import string
import reduction


def get_end_of_chain(
        hash_string: str,
        hash_func: Callable,
        chain_length: int = 1000,
        start_step: int = 0,
        word_length: int = 6,
        alphabet=None
) -> str:
    """
    get_end_of_chain - function designed to get end password in the chain.

    :param hash_string: Input hash string to reduce.
    :param hash_func: Hash function.
    :param chain_length: Max chain length.
    :param start_step: Start step of reduction to generate next chain links.
    :param word_length: Target password length.
    :param chain_length: Max chain length.
    :param alphabet: Input alphabet.
    :return: Last chain link password.
    """
    word: str = ""
    for i in range(start_step, chain_length):
        word = reduction.reduce(
            hash_string=hash_string,
            iteration=i,
            alphabet=alphabet,
            word_length=word_length
        )
        hash_string = hash_func(word)
    return word


def generate(
        quantity: int,
        hash_func: Callable,
        add_to_table: Callable,
        chain_length: int = 1000,
        alphabet=None,
        word_length: int = 6
) -> None:
    """
    generate - function, designed to generate rainbow table.

    :param quantity: Target quantity of table records to generate.
    :param hash_func: Callback for getting hash from the password.
    :param add_to_table: Callback for adding records to rainbow table.
    :param chain_length: Target chain length for one record.
    :param alphabet: Input alphabet.
    :param word_length: Password length.
    """
    if alphabet is None:
        alphabet = list(string.ascii_letters)

    random = Random()
    i = 0
    while i < quantity:
        start_word: str = "".join([alphabet[random.randint(0, len(alphabet) - 1)] for i in range(word_length)])
        end_word: str = get_end_of_chain(
            hash_string=hash_func(start_word),
            hash_func=hash_func,
            chain_length=chain_length,
            word_length=word_length,
            alphabet=alphabet
        )
        if add_to_table(start_word, end_word):
            i += 1


def get_password(
        hash_string: str,
        hash_func: Callable,
        get_from_table: Callable,
        chain_length: int = 1000,
        word_length: int = 6,
        alphabet=None
) -> str:
    """
    get_password - function, designed to calculate password from rainbow table.

    :param chain_length: Target chain length.
    :param get_from_table: Callback fro getting records from table by end_word.
    :param hash_func: Callback for getting hash from the password.
    :param hash_string: Input hash of the password.
    :param word_length: Target password length.
    :param alphabet: Input alphabet.
    :return: Target password
    """
    # Decreasing reduction start.
    for i in range(chain_length - 1, -1, -1):
        # Calculating end of chain for selected reduction start.
        end_word: str = get_end_of_chain(
            hash_string=hash_string,
            hash_func=hash_func,
            chain_length=chain_length,
            start_step=i,
            word_length=word_length,
            alphabet=alphabet
        )
        # Looking for coincidence in the table.
        entries: list = get_from_table(end_word)
        if len(entries):
            # For each coincidence.
            for word in entries:
                # Restoring entire chain.
                for step in range(0, chain_length):
                    current_hash: str = hash_func(word)
                    # If hash matches - congratulations, we have found password!
                    if hash_string == current_hash:
                        return word
                    word = reduction.reduce(
                        hash_string=current_hash,
                        iteration=step,
                        alphabet=alphabet,
                        word_length=word_length
                    )