planetmint/tests/utils.py

# Copyright © 2020 Interplanetary Database Association e.V.,
# Planetmint and IPDB software contributors.
# SPDX-License-Identifier: (Apache-2.0 AND CC-BY-4.0)
# Code is Apache-2.0 and docs are CC-BY-4.0
import multiprocessing
from hashlib import sha3_256

import base58
import base64
import random

from functools import singledispatch

from planetmint import backend
from planetmint.backend.localmongodb.connection import LocalMongoDBConnection
from planetmint.backend.tarantool.sync_io.connection import TarantoolDBConnection
from planetmint.backend.schema import TABLES
from transactions.common import crypto
from transactions.common.transaction_mode_types import BROADCAST_TX_COMMIT
from transactions.types.assets.create import Create
from transactions.types.elections.vote import Vote
from transactions.types.elections.validator_utils import election_id_to_public_key
from planetmint.abci.utils import merkleroot, key_to_base64
from planetmint.abci.rpc import MODE_COMMIT, MODE_LIST


@singledispatch
def flush_db(connection, dbname):
    raise NotImplementedError


@flush_db.register(LocalMongoDBConnection)
def flush_localmongo_db(connection, dbname):
    for t in TABLES:
        getattr(connection.conn[dbname], t).delete_many({})


@flush_db.register(TarantoolDBConnection)
def flush_tarantool_db(connection, dbname):
    connection.connect().call("drop")
    connection.connect().call("init")


def generate_block(planet, test_abci_rpc):
    from transactions.common.crypto import generate_key_pair

    alice = generate_key_pair()
    tx = Create.generate([alice.public_key], [([alice.public_key], 1)], assets=[{"data": None}]).sign(
        [alice.private_key]
    )

    code, message = test_abci_rpc.write_transaction(
        MODE_LIST, test_abci_rpc.tendermint_rpc_endpoint, MODE_COMMIT, tx, BROADCAST_TX_COMMIT
    )
    assert code == 202


def to_inputs(election, i, ed25519_node_keys):
    input0 = election.to_inputs()[i]
    votes = election.outputs[i].amount
    public_key0 = input0.owners_before[0]
    key0 = ed25519_node_keys[public_key0]
    return (input0, votes, key0)


def gen_vote(election, i, ed25519_node_keys):
    (input_i, votes_i, key_i) = to_inputs(election, i, ed25519_node_keys)
    election_pub_key = election_id_to_public_key(election.id)
    return Vote.generate([input_i], [([election_pub_key], votes_i)], election_ids=[election.id]).sign(
        [key_i.private_key]
    )


def generate_validators(powers):
    """Generates an arbitrary number of validators with random public keys.

    The object under the `storage` key is in the format expected by DB.

    The object under the `eleciton` key is in the format expected by
    the upsert validator election.

    `public_key`, `private_key` are in the format used for signing transactions.

    Args:
        powers: A list of intergers representing the voting power to
                assign to the corresponding validators.
    """
    validators = []
    for power in powers:
        kp = crypto.generate_key_pair()
        validators.append(
            {
                "storage": {
                    "public_key": {
                        "value": key_to_base64(base58.b58decode(kp.public_key).hex()),
                        "type": "ed25519-base64",
                    },
                    "voting_power": power,
                },
                "election": {
                    "node_id": f"node-{random.choice(range(100))}",
                    "power": power,
                    "public_key": {
                        "value": base64.b16encode(base58.b58decode(kp.public_key)).decode("utf-8"),
                        "type": "ed25519-base16",
                    },
                },
                "public_key": kp.public_key,
                "private_key": kp.private_key,
            }
        )
    return validators


# NOTE: This works for some but not for all test cases check if this or code base needs fix
def generate_election(b, cls, public_key, private_key, asset_data, voter_keys):
    voters = b.get_recipients_list()
    election = cls.generate([public_key], voters, asset_data, None).sign([private_key])

    votes = [
        Vote.generate([election.to_inputs()[i]], [([election_id_to_public_key(election.id)], power)], [election.id])
        for i, (_, power) in enumerate(voters)
    ]
    for key, v in zip(voter_keys, votes):
        v.sign([key])

    return election, votes


def delete_unspent_outputs(connection, *unspent_outputs):
    """Deletes the given ``unspent_outputs`` (utxos).

    Args:
        *unspent_outputs (:obj:`tuple` of :obj:`dict`): Variable
            length tuple or list of unspent outputs.
    """
    if unspent_outputs:
        return backend.query.delete_unspent_outputs(connection, *unspent_outputs)


def get_utxoset_merkle_root(connection):
    """Returns the merkle root of the utxoset. This implies that
    the utxoset is first put into a merkle tree.

    For now, the merkle tree and its root will be computed each
    time. This obviously is not efficient and a better approach
    that limits the repetition of the same computation when
    unnecesary should be sought. For instance, future optimizations
    could simply re-compute the branches of the tree that were
    affected by a change.

    The transaction hash (id) and output index should be sufficient
    to uniquely identify a utxo, and consequently only that
    information from a utxo record is needed to compute the merkle
    root. Hence, each node of the merkle tree should contain the
    tuple (txid, output_index).

    .. important:: The leaves of the tree will need to be sorted in
        some kind of lexicographical order.

    Returns:
        str: Merkle root in hexadecimal form.
    """
    utxoset = backend.query.get_unspent_outputs(connection)
    # TODO Once ready, use the already pre-computed utxo_hash field.
    # See common/transactions.py for details.
    hashes = [
        sha3_256("{}{}".format(utxo["transaction_id"], utxo["output_index"]).encode()).digest() for utxo in utxoset
    ]
    # TODO Notice the sorted call!
    return merkleroot(sorted(hashes))


def store_unspent_outputs(connection, *unspent_outputs):
    """Store the given ``unspent_outputs`` (utxos).

    Args:
        *unspent_outputs (:obj:`tuple` of :obj:`dict`): Variable
            length tuple or list of unspent outputs.
    """
    if unspent_outputs:
        return backend.query.store_unspent_outputs(connection, *unspent_outputs)


def update_utxoset(connection, transaction):
    """
    Update the UTXO set given ``transaction``. That is, remove
    the outputs that the given ``transaction`` spends, and add the
    outputs that the given ``transaction`` creates.

    Args:
        transaction (:obj:`~planetmint.models.Transaction`): A new
            transaction incoming into the system for which the UTXOF
            set needs to be updated.
    """
    spent_outputs = [spent_output for spent_output in transaction.spent_outputs]
    if spent_outputs:
        delete_unspent_outputs(connection, *spent_outputs)
    store_unspent_outputs(connection, *[utxo._asdict() for utxo in transaction.unspent_outputs])


class ProcessGroup(object):
    def __init__(self, concurrency=None, group=None, target=None, name=None, args=None, kwargs=None, daemon=None):
        self.concurrency = concurrency or multiprocessing.cpu_count()
        self.group = group
        self.target = target
        self.name = name
        self.args = args or ()
        self.kwargs = kwargs or {}
        self.daemon = daemon
        self.processes = []

    def start(self):
        for i in range(self.concurrency):
            proc = multiprocessing.Process(
                group=self.group,
                target=self.target,
                name=self.name,
                args=self.args,
                kwargs=self.kwargs,
                daemon=self.daemon,
            )
            proc.start()
            self.processes.append(proc)