script.py

import os
import sys
import subprocess
import parser
import shutil
import time
import json

DEV_PATH = "/home/ubuntu/lightning_routing/dev"
RPC_PORT_OFFSET = 10000
PEER_PORT_OFFSET = 5000
REST_PORT_OFFSET = 8000


# start bitcoin daemon and directories
def initial_setup(num_nodes):
    # start bitcoin daemon
    btcd = subprocess.Popen(["btcd", "--txindex", "--simnet","--rpcuser=kek", "--rpcpass=kek"])

    # create separate directories
    for i in range(num_nodes):
        if not os.path.exists(os.path.join(DEV_PATH, str(i))):
            os.makedirs(os.path.join(DEV_PATH, str(i)))




# delete derectories and kill all running processes
def teardown(num_nodes):
    for i in range(num_nodes):
        shutil.rmtree(os.path.join(DEV_PATH, str(i)))

    # kill lnd, btcd
    os.system("killall -9 btcd")
    os.system("killall -9 lnd")



# set up the lnd daemons for every port
def setup_lnd_daemons(num_nodes):
    for i in range(num_nodes):
        rpcport = RPC_PORT_OFFSET + i
        listenport = PEER_PORT_OFFSET + i
        restport= REST_PORT_OFFSET + i
        working_dir = os.path.join(DEV_PATH, str(i))

        print rpcport, listenport

        subprocess.Popen(["lnd", "--rpclisten=localhost:" + str(rpcport),\
            "--listen=localhost:" + str(listenport),\
            "--restlisten=localhost:" + str(restport),\
            "--datadir=data", "--logdir=log",\
            "--debuglevel=info", "--bitcoin.simnet", "--bitcoin.active", \
            "--bitcoin.node=btcd", "--btcd.rpcuser=kek", "--btcd.rpcpass=kek", \
            "--noencryptwallet"], cwd=working_dir)



# set up the wallets for every client
def setup_wallets(num_nodes):
    for i in range(num_nodes):
        rpcport = RPC_PORT_OFFSET + i
        working_dir = os.path.join(DEV_PATH, str(i))

        p = subprocess.call(["lncli", "--rpcserver=localhost:" + str(rpcport), \
                "--macaroonpath=data/admin.macaroon", "create"], cwd=working_dir)

        if p != 0:
            print "Error setting up wallet for node", i



# generate bitcoin addresses for each of the nodes
def generate_btc_addresses(num_nodes):
    address_list = []
    for i in range(num_nodes):
        rpcport = RPC_PORT_OFFSET + i
        print rpcport
        working_dir = os.path.join(DEV_PATH, str(i))

        addr_output = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(rpcport), \
                        "--macaroonpath=data/admin.macaroon", "newaddress",\
                        "np2wkh"], cwd=working_dir)

        json_output = json.loads(addr_output)
        address = json_output["address"].strip('\"')
        address_list.append(address)
    return address_list



# find the public key for all of the lnd nodes
# used to setup peer to peer network and to open channels
def find_pub_keys(num_nodes):
    pub_key_list = []
    for i in range(num_nodes):
        rpcport = RPC_PORT_OFFSET + i
        print rpcport
        working_dir = os.path.join(DEV_PATH, str(i))

        info_output = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(rpcport), \
                        "--macaroonpath=data/admin.macaroon", "getinfo"],\
                        cwd=working_dir)

        json_output = json.loads(info_output)
        pub_key_list.append(json_output["identity_pubkey"].strip('\"'))
    return pub_key_list



# set up peer to peer network for all nodes as per the adjacency matrix
# only one direction needs to be set up, so we set it up for src < dest
def setup_peer_to_peer_network(adj_matrix, pub_key_list):
    for i, adj_nodes in adj_matrix.iteritems():
        for j in adj_nodes:
            if i < j:
                peer_port = PEER_PORT_OFFSET + j
                peer_pub_key = pub_key_list[j]

                my_rpcport = RPC_PORT_OFFSET + i
                working_dir = os.path.join(DEV_PATH, str(i))

                p = subprocess.call(["lncli", "--rpcserver=localhost:" + str(my_rpcport), \
                    "--macaroonpath=data/admin.macaroon", "connect",\
                    peer_pub_key + "@localhost:" + str(peer_port)], cwd=working_dir)
                if p != 0:
                    print "Error in setting up link from ", i, " to ", j


            peer_rpcport = RPC_PORT_OFFSET + j
            peer_working_dir = os.path.join(DEV_PATH, str(j))


            peer_op = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(peer_rpcport), \
                "--macaroonpath=data/admin.macaroon", "listpeers"], cwd=peer_working_dir)
            print "list of channels", peer_op



# set up lightning network for all nodes as per the adjacency matrix
# only one direction needs to be set up, so we set it up for src < dest
# for every channel to become valid, we need to also generate 6 blocks each
def setup_lightning_network(adj_matrix, pub_key_list, credit_amt_dict):
    for i, adj_nodes in adj_matrix.iteritems():
        for j in adj_nodes:
            if i < j:
                credit = int(credit_amt_dict[(i, j)])
                peer_pub_key = pub_key_list[j]

                my_rpcport = RPC_PORT_OFFSET + i
                working_dir = os.path.join(DEV_PATH, str(i))

                wallet_bal = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(my_rpcport), \
                    "--macaroonpath=data/admin.macaroon", "walletbalance"], cwd=working_dir)
                print "wallet balance for ", i, " is ", wallet_bal
                print " initiating channel from ", i, " to ", j, " with bal ", credit

                # create a balanced channel so push half the amount to the other side
                p = subprocess.Popen(["lncli", "--rpcserver=localhost:" + str(my_rpcport), \
                    "--macaroonpath=data/admin.macaroon", "openchannel",\
                    "--node_key=" + peer_pub_key, "--local_amt=" + str(credit)], cwd=working_dir)
                    #"--push_amt=" + str(credit/2)], cwd=working_dir)
                time.sleep(5)
                print p
                #if p != 0:
                    #print "Error initiating channel from ", i, " to ", j



                # generate 6 blocks
                p = subprocess.call(["btcctl", "--simnet", "--rpcuser=kek",\
                        "--rpcpass=kek", "generate",  "10"])
                time.sleep(5)
                if p != 0:
                    print "Error in generating blocks for channel confirmation from ", i, " to ", j

                channel_op = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(my_rpcport), \
                    "--macaroonpath=data/admin.macaroon", "pendingchannels"], cwd=working_dir)
                print "list of channels", channel_op



# fund the bitcoin addresses by turning on and off the btcd client
# with each of them being the miner in turn
def fund_btc_addresses(btc_address_list):

    os.system("killall -9 btcd")

    for i, address in enumerate(btc_address_list):
        # kill current bitcoin client and restart with current address
        # receiving mining rewards
        os.system("killall -9 btcd")

        btcd = subprocess.Popen(["btcd", "--txindex", "--simnet","--rpcuser=kek", "--rpcpass=kek",\
                "--miningaddr=" + address])

        time.sleep(5)

        # generate enough blocks for funds
        # TODO: edit this to make sure channel creation doesn't fail
        p = subprocess.call(["btcctl", "--simnet", "--rpcuser=kek",\
                "--rpcpass=kek", "generate",  "400"])
        if p != 0:
            print "error mining blocks for address", address

        time.sleep(5)

    # make sure wallets are funded
    for i, address in enumerate(btc_address_list):

        my_rpcport = RPC_PORT_OFFSET + i
        working_dir = os.path.join(DEV_PATH, str(i))


        wallet_bal = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(my_rpcport), \
            "--macaroonpath=data/admin.macaroon", "walletbalance"], cwd=working_dir)
        print "wallet balance for ", i, " is ", wallet_bal



# execute each transaction and decide if it was a success or a failure
def execute_transactions(transaction_list):
    success_num = 0.0
    success_vol = 0.0
    total_vol = 0.0
    total_num = 0.0

    for src_dest, txn_size in transaction_list.iteritems():

        total_vol += txn_size
        total_num += 1

        src = src_dest[0]
        dst = src_dest[1]

        # generate invoice for payment
        dst_rpcport = RPC_PORT_OFFSET + dst
        dst_dir = os.path.join(DEV_PATH, str(dst))
        invoice_output = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(dst_rpcport), \
                        "--macaroonpath=data/admin.macaroon", "addinvoice", \
                        "--amt=" + str(int(txn_size))], cwd=dst_dir)

        # parse the output to the get the encoded invoice
        json_output = json.loads(invoice_output)
        encoded_invoice = json_output["pay_req"].strip('\"')

        # now make a payment to that invoice
        src_rpcport = RPC_PORT_OFFSET + src
        src_dir = os.path.join(DEV_PATH, str(src))
        payment_output = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(src_rpcport), \
                        "--macaroonpath=data/admin.macaroon", "sendpayment", \
                        "--pay_req=" + encoded_invoice], cwd=src_dir)


        channel_info_output = subprocess.check_output(["lncli", "--rpcserver=localhost:" + str(src_rpcport), \
                        "--macaroonpath=data/admin.macaroon", "listchannels"], cwd=src_dir)

        print channel_info_output



        # TODO: need to parse payment output to infer success/failure
        print payment_output

def main():

    # parse topology to get number of nodes and edge links
    adj_dict = {0: set([1]), 1: set([0])} #parser.parse_graph(sys.argv[1])
    num_nodes = len(adj_dict)

    # read the credit amounts per link
    credit_amt_dict = {(0,1): 20000} # parser.parse_credits(sys.argv[1] + "_CREDIT_LINKS")

    # parse the transactions
    '''txn_srcdest_list, txn_size_list = parser.parse_transactions(sys.argv[2])
    transactions = dict(zip(txn_srcdest_list, txn_size_list))'''

    try:
        initial_setup(num_nodes)

        # create per node directories, run lnd in each of them
        setup_lnd_daemons(num_nodes)
        time.sleep(10)

        # setup wallets per node - not needed because they are unlocked by default
        #setup_wallets(num_nodes)

        # generate bitcoin_addresses for the nodes
        btc_address_list = generate_btc_addresses(num_nodes)
        print "btc address list is", btc_address_list

        # fund each one of them according to necessary credit
        # repeatedly turn on and off btcd which one mining
        fund_btc_addresses(btc_address_list)

        # find public key for every node by parsing output
        pub_key_list = find_pub_keys(num_nodes)
        print(pub_key_list)

        # set up the underlying P2P network
        setup_peer_to_peer_network(adj_dict, pub_key_list)

        time.sleep(10)

        # setup lighting network
        # with the right capacities
        setup_lightning_network(adj_dict, pub_key_list, credit_amt_dict)

        # execute transactions
        '''execute_transactions(transactions)'''

    finally:
        # close everything
        teardown(num_nodes)
        ()


main()