#!/usr/bin/env python3
# Copyright (c) 2017-present The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""
Test how locally submitted transactions are sent to the network when private broadcast is used.
"""

import re
import time
import threading

from test_framework.p2p import (
    P2PDataStore,
    P2PInterface,
    P2P_SERVICES,
    start_p2p_listener,
)
from test_framework.messages import (
    CAddress,
    CInv,
    MSG_WTX,
    malleate_tx_to_invalid_witness,
    msg_inv,
    msg_tx,
)
from test_framework.netutil import (
    format_addr_port
)
from test_framework.script_util import build_malleated_tx_package
from test_framework.socks5 import (
    start_socks5_server,
)
from test_framework.test_framework import (
    BitcoinTestFramework,
)
from test_framework.util import (
    assert_equal,
    assert_greater_than_or_equal,
    assert_not_equal,
    assert_raises_rpc_error,
    tor_port,
)
from test_framework.wallet import (
    MiniWallet,
)

P2P_PRIVATE_VERSION = 70016
NUM_PRIVATE_BROADCAST_PER_TX = 3


class P2PPrivateBroadcast(BitcoinTestFramework):
    def set_test_params(self):
        self.disable_autoconnect = False
        self.num_nodes = 2

    def setup_nodes(self):
        self.destinations = []

        self.destinations_lock = threading.Lock()

        def find_connection_type_in_debug_log(to_addr, to_port):
            """
            Scan the debug log of tx_originator for a connection attempt to to_addr:to_port.
            Return the connection type (outbound-full-relay, private-broadcast, etc) or
            None if there is no connection attempt to to_addr:to_port.
            """
            with open(self.tx_originator_debug_log_path, mode="r", encoding="utf-8") as debug_log:
                for line in debug_log.readlines():
                    match = re.match(f".*trying v. connection \\((.+)\\) to \\[?{to_addr}]?:{to_port},.*", line)
                    if match:
                        return match.group(1)
            return None

        def destinations_factory(requested_to_addr, requested_to_port):
            """
            Instruct the SOCKS5 proxy to redirect connections:
            * The first automatic outbound connection -> P2PDataStore
            * The first private broadcast connection -> nodes[1]
            * Anything else -> P2PInterface
            """
            conn_type = None
            def found_connection_in_debug_log():
                nonlocal conn_type
                conn_type = find_connection_type_in_debug_log(requested_to_addr, requested_to_port)
                return conn_type is not None

            self.wait_until(found_connection_in_debug_log)

            with self.destinations_lock:
                i = len(self.destinations)
                actual_to_addr = ""
                actual_to_port = 0
                listener = None
                target_name = ""
                if conn_type == "private-broadcast" and not any(dest["conn_type"] == "private-broadcast" for dest in self.destinations):
                    # Instruct the SOCKS5 server to redirect the first private
                    # broadcast connection from nodes[0] to nodes[1]
                    actual_to_addr = "127.0.0.1" # nodes[1] listen address
                    actual_to_port = tor_port(1) # nodes[1] listen port for Tor
                    target_name = "nodes[1]"
                else:
                    # Create a Python P2P listening node and instruct the SOCKS5 proxy to
                    # redirect the connection to it. The first outbound connection is used
                    # later to serve GETDATA, thus make it P2PDataStore().
                    if conn_type == "outbound-full-relay" and not any(dest["conn_type"] == "outbound-full-relay" for dest in self.destinations):
                        listener = P2PDataStore()
                        target_name = "Python P2PDataStore"
                    else:
                        listener = P2PInterface()
                        target_name = "Python P2PInterface"
                    listener.peer_connect_helper(dstaddr="0.0.0.0", dstport=0, net=self.chain, timeout_factor=self.options.timeout_factor)
                    listener.peer_connect_send_version(services=P2P_SERVICES)

                    actual_to_addr, actual_to_port = start_p2p_listener(self.network_thread, listener)

                self.log.debug(f"Instructing the SOCKS5 proxy to redirect connection i={i} ({conn_type}) for "
                               f"{format_addr_port(requested_to_addr, requested_to_port)} to "
                               f"{format_addr_port(actual_to_addr, actual_to_port)} ({target_name})")

                self.destinations.append({
                    "requested_to": format_addr_port(requested_to_addr, requested_to_port),
                    "conn_type": conn_type,
                    "node": listener,
                })
                assert_equal(len(self.destinations), i + 1)

                return {
                    "actual_to_addr": actual_to_addr,
                    "actual_to_port": actual_to_port,
                }

        self.socks5_server = start_socks5_server(destinations_factory)

        self.extra_args = [
            [
                # Needed to be able to add CJDNS addresses to addrman (otherwise they are unroutable).
                "-cjdnsreachable",
                # Connecting, sending garbage, being disconnected messes up with this test's
                # check_broadcasts() which waits for a particular Python node to receive a connection.
                "-v2transport=0",
                "-test=addrman",
                "-privatebroadcast",
                f"-proxy={self.socks5_server.conf.addr[0]}:{self.socks5_server.conf.addr[1]}",
                # To increase coverage, make it think that the I2P network is reachable so that it
                # selects such addresses as well. Pick a proxy address where nobody is listening
                # and connection attempts fail quickly.
                "-i2psam=127.0.0.1:1",
            ],
            [
                "-connect=0",
                f"-bind=127.0.0.1:{tor_port(1)}=onion",
            ],
        ]
        super().setup_nodes()

    def setup_network(self):
        self.setup_nodes()

    def check_broadcasts(self, label, tx, broadcasts_to_expect, skip_destinations):
        def wait_and_get_destination(n):
            """Wait for self.destinations[] to have at least n elements and return the 'n'th."""
            def get_destinations_len():
                with self.destinations_lock:
                    return len(self.destinations)
            self.wait_until(lambda: get_destinations_len() > n)
            with self.destinations_lock:
                return self.destinations[n]

        broadcasts_done = 0
        i = skip_destinations - 1
        while broadcasts_done < broadcasts_to_expect:
            i += 1
            self.log.debug(f"{label}: waiting for outbound connection i={i}")
            # At this point the connection may not yet have been established (A),
            # may be active (B), or may have already been closed (C).
            dest = wait_and_get_destination(i)
            peer = dest["node"]
            if peer is None:
                continue # That is the first private broadcast connection, redirected to nodes[1]
            peer.wait_until(lambda: peer.message_count["version"] == 1, check_connected=False)
            # Now it is either (B) or (C).
            if peer.last_message["version"].nServices != 0:
                self.log.debug(f"{label}: outbound connection i={i} to {dest['requested_to']} not a private broadcast, ignoring it (maybe feeler or extra block only)")
                continue
            self.log.debug(f"{label}: outbound connection i={i} to {dest['requested_to']} must be a private broadcast, checking it")
            peer.wait_for_disconnect()
            # Now it is (C).
            assert_equal(peer.message_count, {
                "version": 1,
                "verack": 1,
                "inv": 1,
                "tx": 1,
                "ping": 1
            })
            dummy_address = CAddress()
            dummy_address.nServices = 0
            assert_equal(peer.last_message["version"].nVersion, P2P_PRIVATE_VERSION)
            assert_equal(peer.last_message["version"].nServices, 0)
            assert_equal(peer.last_message["version"].nTime, 0)
            assert_equal(peer.last_message["version"].addrTo, dummy_address)
            assert_equal(peer.last_message["version"].addrFrom, dummy_address)
            assert_equal(peer.last_message["version"].strSubVer, "/pynode:0.0.1/")
            assert_equal(peer.last_message["version"].nStartingHeight, 0)
            assert_equal(peer.last_message["version"].relay, 0)
            assert_equal(peer.last_message["tx"].tx.txid_hex, tx["txid"])
            self.log.info(f"{label}: ok: outbound connection i={i} is private broadcast of txid={tx['txid']}")
            broadcasts_done += 1

        # Verify the tx we just observed is tracked in getprivatebroadcastinfo.
        pbinfo = self.nodes[0].getprivatebroadcastinfo()
        pending = [t for t in pbinfo["transactions"] if t["txid"] == tx["txid"] and t["wtxid"] == tx["wtxid"]]
        assert_equal(len(pending), 1)
        assert_equal(pending[0]["hex"].lower(), tx["hex"].lower())
        peers = pending[0]["peers"]
        assert len(peers) >= NUM_PRIVATE_BROADCAST_PER_TX
        assert all("address" in p and "sent" in p for p in peers)
        assert_greater_than_or_equal(sum(1 for p in peers if "received" in p), broadcasts_to_expect)

    def run_test(self):
        tx_originator = self.nodes[0]
        self.tx_originator_debug_log_path = tx_originator.debug_log_path
        tx_receiver = self.nodes[1]
        far_observer = tx_receiver.add_p2p_connection(P2PInterface())

        self.log.info("Test getprivatebroadcastinfo and abortprivatebroadcast fails if the node is running without -privatebroadcast set")
        assert_raises_rpc_error(-32601, "Private broadcast is not enabled. Ensure you're running Bitcoin Core with -privatebroadcast=1.",
            tx_receiver.getprivatebroadcastinfo)
        assert_raises_rpc_error(-32601, "Private broadcast is not enabled. Ensure you're running Bitcoin Core with -privatebroadcast=1.",
            tx_receiver.abortprivatebroadcast, "00" * 32)

        self.fill_node_addrman(node_index=0, address_types_to_add=[CAddress.NET_IPV4, CAddress.NET_IPV6, CAddress.NET_TORV3, CAddress.NET_I2P, CAddress.NET_CJDNS])

        wallet = MiniWallet(tx_originator)

        txs = wallet.create_self_transfer_chain(chain_length=3)
        self.log.info(f"Created txid={txs[0]['txid']}: for basic test")
        self.log.info(f"Created txid={txs[1]['txid']}: for broadcast with dependency in mempool + rebroadcast")
        self.log.info(f"Created txid={txs[2]['txid']}: for broadcast with dependency not in mempool")
        tx_originator.sendrawtransaction(hexstring=txs[0]["hex"], maxfeerate=0.1)

        self.log.info("First private broadcast: waiting for the transaction to reach the recipient")
        self.wait_until(lambda: len(tx_receiver.getrawmempool()) > 0)
        self.log.info("First private broadcast: the recipient received the transaction")
        far_observer.wait_for_tx(txs[0]["txid"])
        self.log.info("First private broadcast: the recipient further relayed the transaction")

        # One already checked above, check the other NUM_PRIVATE_BROADCAST_PER_TX - 1 broadcasts.
        self.check_broadcasts("Basic", txs[0], NUM_PRIVATE_BROADCAST_PER_TX - 1, 0)

        self.log.info("Resending the same transaction via RPC again (it is not in the mempool yet)")
        ignoring_msg = f"Ignoring unnecessary request to schedule an already scheduled transaction: txid={txs[0]['txid']}, wtxid={txs[0]['wtxid']}"
        with tx_originator.busy_wait_for_debug_log(expected_msgs=[ignoring_msg.encode()]):
            tx_originator.sendrawtransaction(hexstring=txs[0]["hex"], maxfeerate=0)

        self.log.info("Sending a malleated transaction with an invalid witness via RPC")
        malleated_invalid = malleate_tx_to_invalid_witness(txs[0])
        assert_raises_rpc_error(-26, "mempool-script-verify-flag-failed",
                                tx_originator.sendrawtransaction,
                                hexstring=malleated_invalid.serialize_with_witness().hex(),
                                maxfeerate=0.1)

        self.log.info("Checking that the transaction is not in the originator node's mempool")
        assert_equal(len(tx_originator.getrawmempool()), 0)

        wtxid_int = int(txs[0]["wtxid"], 16)
        inv = CInv(MSG_WTX, wtxid_int)

        tx_returner = None # First outbound-full-relay, will be P2PDataStore.
        other_peer = None # Any other outbound-full-relay, we use the second one.

        def set_tx_returner_and_other():
            nonlocal tx_returner
            nonlocal other_peer
            tx_returner = None
            other_peer = None
            with self.destinations_lock:
                for dest in self.destinations:
                    if dest["conn_type"] == "outbound-full-relay" and dest["node"] is not None:
                        if tx_returner is None:
                            assert(type(dest["node"]) is P2PDataStore)
                            tx_returner = dest["node"]
                        else:
                            assert(type(dest["node"]) is P2PInterface)
                            other_peer = dest["node"]
                            return True
            return False

        self.wait_until(set_tx_returner_and_other)

        tx_returner.wait_for_connect()
        other_peer.wait_for_connect()

        self.log.info("Sending INV and waiting for GETDATA from node")
        tx_returner.tx_store[wtxid_int] = txs[0]["tx"]
        assert "getdata" not in tx_returner.last_message
        received_back_msg = f"Received our privately broadcast transaction (txid={txs[0]['txid']}) from the network"
        with tx_originator.assert_debug_log(expected_msgs=[received_back_msg]):
            tx_returner.send_without_ping(msg_inv([inv]))
            tx_returner.wait_until(lambda: "getdata" in tx_returner.last_message)
            self.wait_until(lambda: len(tx_originator.getrawmempool()) > 0)

        self.log.info("Waiting for normal broadcast to another peer")
        other_peer.wait_for_inv([inv])

        self.log.info("Checking getprivatebroadcastinfo no longer reports the transaction after it is received back")
        pbinfo = tx_originator.getprivatebroadcastinfo()
        pending = [t for t in pbinfo["transactions"] if t["txid"] == txs[0]["txid"] and t["wtxid"] == txs[0]["wtxid"]]
        assert_equal(len(pending), 0)

        self.log.info("Sending a transaction that is already in the mempool")
        skip_destinations = len(self.destinations)
        tx_originator.sendrawtransaction(hexstring=txs[0]["hex"], maxfeerate=0)
        self.check_broadcasts("Broadcast of mempool transaction", txs[0], NUM_PRIVATE_BROADCAST_PER_TX, skip_destinations)

        self.log.info("Sending a transaction with a dependency in the mempool")
        skip_destinations = len(self.destinations)
        tx_originator.sendrawtransaction(hexstring=txs[1]["hex"], maxfeerate=0.1)
        self.check_broadcasts("Dependency in mempool", txs[1], NUM_PRIVATE_BROADCAST_PER_TX, skip_destinations)

        self.log.info("Sending a transaction with a dependency not in the mempool (should be rejected)")
        assert_equal(len(tx_originator.getrawmempool()), 1)
        assert_raises_rpc_error(-25, "bad-txns-inputs-missingorspent",
                                tx_originator.sendrawtransaction, hexstring=txs[2]["hex"], maxfeerate=0.1)
        assert_raises_rpc_error(-25, "bad-txns-inputs-missingorspent",
                                tx_originator.sendrawtransaction, hexstring=txs[2]["hex"], maxfeerate=0)

        # Since txs[1] has not been received back by tx_originator,
        # it should be re-broadcast after a while. Advance tx_originator's clock
        # to trigger a re-broadcast. Should be more than the maximum returned by
        # NextTxBroadcast() in net_processing.cpp.
        self.log.info("Checking that rebroadcast works")
        delta = 20 * 60 # 20min
        skip_destinations = len(self.destinations)
        rebroadcast_msg = f"Reattempting broadcast of stale txid={txs[1]['txid']}"
        with tx_originator.busy_wait_for_debug_log(expected_msgs=[rebroadcast_msg.encode()]):
            tx_originator.setmocktime(int(time.time()) + delta)
            tx_originator.mockscheduler(delta)
        self.check_broadcasts("Rebroadcast", txs[1], 1, skip_destinations)
        tx_originator.setmocktime(0) # Let the clock tick again (it will go backwards due to this).

        self.log.info("Sending a pair of transactions with the same txid but different valid wtxids via RPC")
        parent = wallet.create_self_transfer()["tx"]
        parent_amount = parent.vout[0].nValue - 10000
        child_amount = parent_amount - 10000
        siblings_parent, sibling1, sibling2 = build_malleated_tx_package(
            parent=parent,
            rebalance_parent_output_amount=parent_amount,
            child_amount=child_amount)
        self.log.info(f"  - sibling1: txid={sibling1.txid_hex}, wtxid={sibling1.wtxid_hex}")
        self.log.info(f"  - sibling2: txid={sibling2.txid_hex}, wtxid={sibling2.wtxid_hex}")
        assert_equal(sibling1.txid_hex, sibling2.txid_hex)
        assert_not_equal(sibling1.wtxid_hex, sibling2.wtxid_hex)
        assert_equal(len(tx_originator.getrawmempool()), 1)
        tx_returner.send_without_ping(msg_tx(siblings_parent))
        self.wait_until(lambda: len(tx_originator.getrawmempool()) > 1)
        self.log.info("  - siblings' parent added to the mempool")
        tx_originator.sendrawtransaction(hexstring=sibling1.serialize_with_witness().hex(), maxfeerate=0.1)
        self.log.info("  - sent sibling1: ok")
        tx_originator.sendrawtransaction(hexstring=sibling2.serialize_with_witness().hex(), maxfeerate=0.1)
        self.log.info("  - sent sibling2: ok")

        self.log.info("Checking abortprivatebroadcast removes a pending private-broadcast transaction")
        tx_abort = wallet.create_self_transfer()
        tx_originator.sendrawtransaction(hexstring=tx_abort["hex"], maxfeerate=0.1)
        assert tx_abort["wtxid"] in [t["wtxid"] for t in tx_originator.getprivatebroadcastinfo()["transactions"]]
        abort_res = tx_originator.abortprivatebroadcast(tx_abort["txid"])
        assert_equal(len(abort_res["removed_transactions"]), 1)
        assert_equal(abort_res["removed_transactions"][0]["txid"], tx_abort["txid"])
        assert_equal(abort_res["removed_transactions"][0]["wtxid"], tx_abort["wtxid"])
        assert_equal(abort_res["removed_transactions"][0]["hex"].lower(), tx_abort["hex"].lower())
        assert all(t["wtxid"] != tx_abort["wtxid"] for t in tx_originator.getprivatebroadcastinfo()["transactions"])

        self.log.info("Checking abortprivatebroadcast fails for non-existent transaction")
        assert_raises_rpc_error(
            -5,
            "Transaction not in private broadcast queue",
            tx_originator.abortprivatebroadcast,
            "0" * 64,
        )

        # Stop the SOCKS5 proxy server to avoid it being upset by the bitcoin
        # node disconnecting in the middle of the SOCKS5 handshake when we
        # restart below.
        self.socks5_server.stop()

        self.log.info("Trying to send a transaction when none of Tor or I2P is reachable")
        self.restart_node(0, extra_args=[
            "-privatebroadcast",
            "-v2transport=0",
            # A location where definitely a Tor control is not listening. This would allow
            # Bitcoin Core to start, hoping/assuming that the location of the Tor proxy
            # may be retrieved after startup from the Tor control, but it will not be, so
            # the RPC should throw.
            "-torcontrol=127.0.0.1:1",
            "-listenonion",
        ])
        assert_raises_rpc_error(-1, "none of the Tor or I2P networks is reachable",
                                tx_originator.sendrawtransaction, hexstring=txs[0]["hex"], maxfeerate=0.1)


if __name__ == "__main__":
    P2PPrivateBroadcast(__file__).main()