From e11418871f7503e227e34bb4726293b97806e110 Mon Sep 17 00:00:00 2001
From: mikera <mike@mikera.net>
Date: Tue, 6 Aug 2024 14:54:23 +0100
Subject: [PATCH] Improve CLI for transactions and queries

---
 .../src/main/java/convex/cli/ACommand.java    |   6 +-
 .../convex/cli/client/AClientCommand.java     | 100 ++++++++++--------
 .../main/java/convex/cli/client/Query.java    |   6 +-
 .../main/java/convex/cli/client/Transact.java |  14 +--
 .../java/convex/cli/mixins/AddressMixin.java  |  16 +--
 .../java/convex/cli/mixins/StoreMixin.java    |  66 ++++++++----
 6 files changed, 120 insertions(+), 88 deletions(-)

diff --git a/convex-cli/src/main/java/convex/cli/ACommand.java b/convex-cli/src/main/java/convex/cli/ACommand.java
index d0384bc62..2c2489e06 100644
--- a/convex-cli/src/main/java/convex/cli/ACommand.java
+++ b/convex-cli/src/main/java/convex/cli/ACommand.java
@@ -2,7 +2,6 @@
 
 import java.io.Console;
 import java.io.IOException;
-import java.util.Scanner;
 
 import convex.cli.output.Coloured;
 import convex.cli.output.RecordOutput;
@@ -96,10 +95,7 @@ public String prompt(String message) {
 		
 		if (isColoured()) message=Coloured.blue(message);
 		inform(0,message);
-		try (Scanner scanner = new Scanner(System.in)) {
-			String s=scanner.nextLine();
-			return s;
-		}
+		return System.console().readLine();
 	}
 	
 	public char[] readPassword(String prompt) {
diff --git a/convex-cli/src/main/java/convex/cli/client/AClientCommand.java b/convex-cli/src/main/java/convex/cli/client/AClientCommand.java
index 6cfbf87f4..ca650cbba 100644
--- a/convex-cli/src/main/java/convex/cli/client/AClientCommand.java
+++ b/convex-cli/src/main/java/convex/cli/client/AClientCommand.java
@@ -1,22 +1,16 @@
 package convex.cli.client;
 
-import java.util.concurrent.TimeUnit;
-
 import convex.api.Convex;
 import convex.cli.ATopCommand;
 import convex.cli.CLIError;
-import convex.cli.Constants;
+import convex.cli.ExitCodes;
 import convex.cli.mixins.AddressMixin;
 import convex.cli.mixins.KeyMixin;
 import convex.cli.mixins.RemotePeerMixin;
 import convex.cli.mixins.StoreMixin;
-import convex.core.Result;
 import convex.core.crypto.AKeyPair;
-import convex.core.data.ABlob;
-import convex.core.data.ACell;
+import convex.core.data.AccountKey;
 import convex.core.data.Address;
-import convex.core.lang.Symbols;
-import convex.core.lang.ops.Special;
 import picocli.CommandLine.Mixin;
 import picocli.CommandLine.Option;
 
@@ -36,7 +30,7 @@ public abstract class AClientCommand extends ATopCommand {
 	
 	@Option(names={"--timeout"},
 			description="Timeout in miliseconds.")
-	protected long timeout = Constants.DEFAULT_TIMEOUT_MILLIS;
+	protected Long timeout;
 
 
 	/**
@@ -45,64 +39,82 @@ public abstract class AClientCommand extends ATopCommand {
 	 */
 	protected Convex clientConnect() {
 		try {
-		return peerMixin.connect();
+			Convex convex= peerMixin.connect();
+			if (timeout!=null) {
+				convex.setTimeout(timeout);
+			}
+			return convex;
 		} catch (Exception ex) {
 			throw new CLIError("Unable to connect to Convex: "+ex.getMessage(),ex);
 		}
 	}
 	
+	/**
+	 * Connect to Convex ready to query
+	 * @return
+	 */
+	protected Convex connectQuery() {
+		Convex convex=clientConnect();
+		Address a=getUserAddress();
+		convex.setAddress(a);
+		return convex;
+	}
 	
-	public Address getUserAddress() {
-		Address result= addressMixin.getAddress("Enter Convex account address: ");	
-		return result;
+	/**
+	 * Connect to Convex ready to transact
+	 * @return
+	 */
+	protected Convex connectTransact() {
+		Convex convex=connectQuery();
+		ensureKeyPair(convex);
+		return convex;
 	}
 	
-	protected boolean ensureAddress(Convex convex) {
-		Address a = getUserAddress();
-		if (a!=null) {
-			convex.setAddress(a);
-			return true;
-		}
-		return false;
+	/**
+	 * Gets user address, prompting of not provided.
+	 * @return Valid Address or null if Address not valid
+	 */
+	public Address getUserAddress() {
+		return addressMixin.getAddress("Enter Convex user account address: ");	
 	}
 	
-	protected boolean ensureKeyPair(Convex convex) {
+	
+	protected void ensureKeyPair(Convex convex) {
+		Address a=convex.getAddress();
 		AKeyPair keyPair = convex.getKeyPair();
-		if (keyPair!=null) return true;
+		if (keyPair!=null) return;
 
-		Address address=convex.getAddress();
-		
-		// Try to identify the required keypair for the Address
-		Result ar;
-		try {
-			ar = convex.query(Special.forSymbol(Symbols.STAR_KEY)).get(1000,TimeUnit.MILLISECONDS);
-		} catch (Exception e) {
-			ar=null;
-		}
-		
-		String pk=null;
-		if (ar==null) {
-			// we couldn't query the *key*, so prompt the user
-		} else if (!ar.isError()) {
-			// Attempt to use query result as public key
-			ACell v=ar.getValue();
-			if (v instanceof ABlob) {
-				pk=((ABlob)v).toHexString();
+		String pk=keyMixin.getPublicKey();
+		if (pk==null) {
+			paranoia("You must set --key explicitly in strict security mode");
+			
+			AccountKey k=convex.getAccountKey(a);
+			if (k!=null) {
+				pk=k.toHexString();
+				inform("Address "+a+" requires public key "+pk);
+			} else if (isInteractive()) {
+				pk=prompt("Enter public key for Address "+a+": ");
+			} else {
+				throw new CLIError(ExitCodes.USAGE,"Public key required.");
 			}
 		}
 		
-		if (pk==null) {
-			pk=prompt("Enter public key: ");
+		storeMixin.loadKeyStore();
+		int c=storeMixin.keyCount(pk);
+		if (c==0) {
+			throw new CLIError(ExitCodes.CONFIG,"Can't find keypair with public key: "+pk);
+		} else if (c>1) {
+			throw new CLIError(ExitCodes.CONFIG,"Multiple key pairs found");
 		}
 		
 		keyPair=storeMixin.loadKeyFromStore(pk,keyMixin.getKeyPassword());
 		if (keyPair==null) {
 			// We didn't find required keypair
-			throw new CLIError("Can't find keypair with public key "+pk+" for Address "+address);
+			throw new CLIError(ExitCodes.CONFIG,"Can't find keypair with public key: "+pk);
 		}
 		convex.setKeyPair(keyPair);
-		return true;
 	}
 
 
+
 }
diff --git a/convex-cli/src/main/java/convex/cli/client/Query.java b/convex-cli/src/main/java/convex/cli/client/Query.java
index 13ffe5cb0..6305ad55b 100644
--- a/convex-cli/src/main/java/convex/cli/client/Query.java
+++ b/convex-cli/src/main/java/convex/cli/client/Query.java
@@ -25,7 +25,7 @@ public class Query extends AClientCommand {
 
 	@Parameters(
 			paramLabel="queryCommand", 
-			description="Query command(s). Multiple commands will be executed in sequence unless one fails")
+			description="Query command(s). Multiple commands will be executed in sequence unless one fails.")
 	private String[] commands;
 
 	@Override
@@ -37,11 +37,11 @@ public void run() {
 		}
 
 		try {
-			Convex convex =  clientConnect();
+			Convex convex =  connectQuery();
 			for (int i=0; i<commands.length; i++) {
 				ACell message = Reader.read(commands[i]);
 				Result result = convex.querySync(message, timeout);
-				cli().printResult(result);
+				printResult(result);
 				if (result.isError()) {
 					break;
 				}
diff --git a/convex-cli/src/main/java/convex/cli/client/Transact.java b/convex-cli/src/main/java/convex/cli/client/Transact.java
index 0be1744ac..61aab6614 100644
--- a/convex-cli/src/main/java/convex/cli/client/Transact.java
+++ b/convex-cli/src/main/java/convex/cli/client/Transact.java
@@ -39,15 +39,11 @@ public class Transact extends AClientCommand {
 	@Override
 	public void run() {
 
-		Convex convex = clientConnect();
-		if (!ensureAddress(convex)) {	
-			throw new CLIError(ExitCodes.USAGE,"Must specify a valid address for transaction.");
-		}
-			
-		if (!ensureKeyPair(convex)) {	
-			throw new CLIError(ExitCodes.USAGE,"Must provide a key pair to sign transaction.");
-		}
-
+		Address a=getUserAddress();
+		if (a==null) throw new CLIError(ExitCodes.USAGE,"You must specify a valid origin address for the transaction.");
+		
+		Convex convex = connectTransact();
+		
 		Address address=convex.getAddress();
 		log.trace("Executing transaction: '{}'\n", transactionCode);
 			
diff --git a/convex-cli/src/main/java/convex/cli/mixins/AddressMixin.java b/convex-cli/src/main/java/convex/cli/mixins/AddressMixin.java
index 26f0d5533..9c6a07ab3 100644
--- a/convex-cli/src/main/java/convex/cli/mixins/AddressMixin.java
+++ b/convex-cli/src/main/java/convex/cli/mixins/AddressMixin.java
@@ -8,10 +8,17 @@ public class AddressMixin extends AMixin {
 	
 	@Option(names={"-a", "--address"},
 			defaultValue="${env:CONVEX_ADDRESS}",
-			description = "Account address to use. Can specify with CONVEX_ADDRESS environment variable. with Defaulting to: ${DEFAULT-VALUE}")
+			description = "Account address to use. Can specify with CONVEX_ADDRESS environment variable.}")
 	protected String addressValue = null;
 
+	private Address address=null;
+	
+	/*
+	 * Get user address
+	 */
 	public Address getAddress(String prompt) {
+		if (address!=null) return address;
+		
 		if (addressValue==null) {
 			if ((prompt!=null)&&isInteractive()) {
 				addressValue=prompt(prompt);
@@ -20,11 +27,8 @@ public Address getAddress(String prompt) {
 			}
 		} 
 		
-		Address a = Address.parse(addressValue);
-		if (a==null) {
-			throw new CLIError("Unable to parse --address argument. Should be a numerical address like '#789'. '#' is optional.");
-		}
-		return a;
+		address = Address.parse(addressValue);
+		return address;
 	}
 	
 
diff --git a/convex-cli/src/main/java/convex/cli/mixins/StoreMixin.java b/convex-cli/src/main/java/convex/cli/mixins/StoreMixin.java
index 8f6b21e71..62cc7442a 100644
--- a/convex-cli/src/main/java/convex/cli/mixins/StoreMixin.java
+++ b/convex-cli/src/main/java/convex/cli/mixins/StoreMixin.java
@@ -124,7 +124,7 @@ private KeyStore loadKeyStore(boolean shouldCreate, char[] password) {
 			if (keyFile.exists()) {
 				keyStore = PFXTools.loadStore(keyFile, password);
 			} else if (shouldCreate) {
-				log.warn("No keystore exists, creating at: " + keyFile.getCanonicalPath());
+				informWarning("No keystore exists, creating at: " + keyFile.getCanonicalPath());
 				Utils.ensurePath(keyFile);
 				keyStore = PFXTools.createStore(keyFile, password);
 			} else {
@@ -154,8 +154,8 @@ public void saveKeyStore(char[] storePassword) {
 			throw new CLIError("Trying to save a keystore that has not been loaded!");
 		try {
 			PFXTools.saveStore(keyStore, getKeyStoreFile(), storePassword);
-		} catch (Throwable t) {
-			throw Utils.sneakyThrow(t);
+		} catch (Exception t) {
+			throw new CLIError("Failed to save keystore",t);
 		}
 	}
 
@@ -170,16 +170,26 @@ public void addKeyPairToStore(AKeyPair keyPair, char[] keyPassword) {
 	
 		KeyStore keyStore = getKeystore();
 		if (keyStore == null) {
-			throw new CLIError("Trying to add key pair but keystore is not loaded");
+			throw new CLIError("Trying to add key pair but keystore is not yet loaded!");
 		}
 		try {
 			// save the key in the keystore
 			PFXTools.setKeyPair(keyStore, keyPair, keyPassword);
-		} catch (Throwable t) {
-			throw new CLIError("Cannot store the key to the key store " + t);
+		} catch (Exception t) {
+			throw new CLIError("Cannot store the key to the key store",t);
 		}
 	
 	}
+	
+	public static String trimKey(String publicKey) {
+		publicKey = publicKey.trim();
+
+		publicKey = publicKey.toLowerCase().replaceFirst("^0x", "").strip();
+		if (publicKey.isEmpty()) {
+			return null;
+		}
+		return publicKey;
+	}
 
 	/**
 	 * Loads a keypair from configured keystore
@@ -191,23 +201,13 @@ public AKeyPair loadKeyFromStore(String publicKey, char[] keyPassword) {
 		if (publicKey == null)
 			return null;
 	
-		publicKey = publicKey.trim();
-		publicKey = publicKey.toLowerCase().replaceFirst("^0x", "").strip();
-		if (publicKey.isEmpty()) {
-			return null;
-		}
-	
-		char[] storePassword = getStorePassword();
-	
-		File keyFile = getKeyStoreFile();
+		publicKey = trimKey(publicKey);
+		if (publicKey==null) return null;		
+
 		try {
-			if (!keyFile.exists()) {
-				throw new CLIError("Cannot find keystore file " + keyFile.getCanonicalPath());
-			}
-			KeyStore keyStore = PFXTools.loadStore(keyFile, storePassword);
+			KeyStore keyStore = ensureKeyStore();
 	
 			Enumeration<String> aliases = keyStore.aliases();
-	
 			while (aliases.hasMoreElements()) {
 				String alias = aliases.nextElement();
 				if (alias.indexOf(publicKey) == 0) {
@@ -217,8 +217,32 @@ public AKeyPair loadKeyFromStore(String publicKey, char[] keyPassword) {
 			}
 			return null;
 		} catch (Exception t) {
-			throw new CLIError("Cannot load key store", t);
+			throw new CLIError("Cannot load aliases from key Store", t);
+		}
+	}
+
+	public boolean hasSingleKey(String pk) {
+		return keyCount(pk)==1;
+	}
+
+	public int keyCount(String pk) {
+		pk=trimKey(pk);
+		
+		int result=0;
+		try {
+			KeyStore keyStore = ensureKeyStore();
+	
+			Enumeration<String> aliases = keyStore.aliases();
+			while (aliases.hasMoreElements()) {
+				String alias = aliases.nextElement();
+				if (alias.indexOf(pk) == 0) {
+					result++;
+				}
+			}
+		} catch (Exception t) {
+			throw new CLIError("Cannot load aliases from key Store", t);
 		}
+		return result;
 	}
 
 }