From c7bed34bcaf59fe180c7923e9ed15ea58a315fca Mon Sep 17 00:00:00 2001
From: Wilmer Paulino <wilmer.paulino@gmail.com>
Date: Fri, 19 Feb 2021 17:41:56 -0800
Subject: [PATCH] lncli: add account flag to wallet related commands

---
 cmd/lncli/commands.go         | 22 ++++++++++++++++++----
 cmd/lncli/walletrpc_active.go | 17 +++++++++++++++--
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/cmd/lncli/commands.go b/cmd/lncli/commands.go
index 51d43eae..cde9e9fd 100644
--- a/cmd/lncli/commands.go
+++ b/cmd/lncli/commands.go
@@ -127,6 +127,13 @@ var newAddressCommand = cli.Command{
 	Category:  "Wallet",
 	Usage:     "Generates a new address.",
 	ArgsUsage: "address-type",
+	Flags: []cli.Flag{
+		cli.StringFlag{
+			Name: "account",
+			Usage: "(optional) the name of the account to " +
+				"generate a new address for",
+		},
+	},
 	Description: `
 	Generate a wallet new address. Address-types has to be one of:
 	    - p2wkh:  Pay to witness key hash
@@ -136,13 +143,16 @@ var newAddressCommand = cli.Command{
 
 func newAddress(ctx *cli.Context) error {
 	ctxc := getContext()
-	client, cleanUp := getClient(ctx)
-	defer cleanUp()
 
-	stringAddrType := ctx.Args().First()
+	// Display the command's help message if we do not have the expected
+	// number of arguments/flags.
+	if ctx.NArg() != 1 || ctx.NumFlags() > 1 {
+		return cli.ShowCommandHelp(ctx, "newaddress")
+	}
 
 	// Map the string encoded address type, to the concrete typed address
 	// type enum. An unrecognized address type will result in an error.
+	stringAddrType := ctx.Args().First()
 	var addrType lnrpc.AddressType
 	switch stringAddrType { // TODO(roasbeef): make them ints on the cli?
 	case "p2wkh":
@@ -154,8 +164,12 @@ func newAddress(ctx *cli.Context) error {
 			"are: p2wkh and np2wkh", stringAddrType)
 	}
 
+	client, cleanUp := getClient(ctx)
+	defer cleanUp()
+
 	addr, err := client.NewAddress(ctxc, &lnrpc.NewAddressRequest{
-		Type: addrType,
+		Type:    addrType,
+		Account: ctx.String("account"),
 	})
 	if err != nil {
 		return err
diff --git a/cmd/lncli/walletrpc_active.go b/cmd/lncli/walletrpc_active.go
index 0f13f58d..dfc17672 100644
--- a/cmd/lncli/walletrpc_active.go
+++ b/cmd/lncli/walletrpc_active.go
@@ -523,6 +523,11 @@ var fundPsbtCommand = cli.Command{
 			Usage: "a manual fee expressed in sat/vbyte that " +
 				"should be used when creating the transaction",
 		},
+		cli.StringFlag{
+			Name: "account",
+			Usage: "(optional) the name of the account to use to " +
+				"create/fund the PSBT",
+		},
 	},
 	Action: actionDecorator(fundPsbt),
 }
@@ -536,7 +541,9 @@ func fundPsbt(ctx *cli.Context) error {
 		return cli.ShowCommandHelp(ctx, "fund")
 	}
 
-	req := &walletrpc.FundPsbtRequest{}
+	req := &walletrpc.FundPsbtRequest{
+		Account: ctx.String("account"),
+	}
 
 	// Parse template flags.
 	switch {
@@ -691,6 +698,11 @@ var finalizePsbtCommand = cli.Command{
 			Name:  "funded_psbt",
 			Usage: "the base64 encoded PSBT to finalize",
 		},
+		cli.StringFlag{
+			Name: "account",
+			Usage: "(optional) the name of the account to " +
+				"finalize the PSBT with",
+		},
 	},
 	Action: actionDecorator(finalizePsbt),
 }
@@ -700,7 +712,7 @@ func finalizePsbt(ctx *cli.Context) error {
 
 	// Display the command's help message if we do not have the expected
 	// number of arguments/flags.
-	if ctx.NArg() != 1 && ctx.NumFlags() != 1 {
+	if ctx.NArg() > 1 || ctx.NumFlags() > 2 {
 		return cli.ShowCommandHelp(ctx, "finalize")
 	}
 
@@ -723,6 +735,7 @@ func finalizePsbt(ctx *cli.Context) error {
 	}
 	req := &walletrpc.FinalizePsbtRequest{
 		FundedPsbt: psbtBytes,
+		Account:    ctx.String("account"),
 	}
 
 	walletClient, cleanUp := getWalletClient(ctx)