Coverage Summary for Class: Wallet (co.rsk.core)

Class Class, % Method, % Line, %
Wallet 100% (1/1) 94.4% (17/18) 85.3% (99/116)


1 /* 2  * This file is part of RskJ 3  * Copyright (C) 2017 RSK Labs Ltd. 4  * 5  * This program is free software: you can redistribute it and/or modify 6  * it under the terms of the GNU Lesser General Public License as published by 7  * the Free Software Foundation, either version 3 of the License, or 8  * (at your option) any later version. 9  * 10  * This program is distributed in the hope that it will be useful, 11  * but WITHOUT ANY WARRANTY; without even the implied warranty of 12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13  * GNU Lesser General Public License for more details. 14  * 15  * You should have received a copy of the GNU Lesser General Public License 16  * along with this program. If not, see <http://www.gnu.org/licenses/>. 17  */ 18  19 package co.rsk.core; 20  21 import co.rsk.bitcoinj.core.Sha256Hash; 22 import co.rsk.crypto.EncryptedData; 23 import co.rsk.crypto.KeyCrypterAes; 24 import org.ethereum.core.Account; 25 import org.ethereum.crypto.ECKey; 26 import org.ethereum.crypto.Keccak256Helper; 27 import org.ethereum.datasource.KeyValueDataSource; 28 import org.ethereum.rpc.TypeConverter; 29 import org.bouncycastle.crypto.params.KeyParameter; 30  31 import javax.annotation.concurrent.GuardedBy; 32 import java.io.*; 33 import java.nio.charset.StandardCharsets; 34 import java.util.*; 35  36 public class Wallet { 37  @GuardedBy("accessLock") 38  private final KeyValueDataSource keyDS; 39  40  @GuardedBy("accessLock") 41  private final Map<RskAddress, byte[]> accounts = new HashMap<>(); 42  43  @GuardedBy("accessLock") 44  private final List<RskAddress> initialAccounts = new ArrayList<>(); 45  46  private final Object accessLock = new Object(); 47  private final Map<RskAddress, Long> unlocksTimeouts = new HashMap<>(); 48  49  public Wallet(KeyValueDataSource keyDS) { 50  this.keyDS = keyDS; 51  } 52  53  public List<byte[]> getAccountAddresses() { 54  List<byte[]> addresses = new ArrayList<>(); 55  Set<RskAddress> keys = new HashSet<>(); 56  57  synchronized(accessLock) { 58  for (RskAddress addr: this.initialAccounts) { 59  addresses.add(addr.getBytes()); 60  } 61  62  for (byte[] address: keyDS.keys()) { 63  keys.add(new RskAddress(address)); 64  } 65  66  keys.addAll(accounts.keySet()); 67  keys.removeAll(this.initialAccounts); 68  69  for (RskAddress addr: keys) { 70  addresses.add(addr.getBytes()); 71  } 72  } 73  74  return addresses; 75  } 76  77  public String[] getAccountAddressesAsHex() { 78  return getAccountAddresses().stream() 79  .map(TypeConverter::toJsonHex) 80  .toArray(String[]::new); 81  } 82  83  public RskAddress addAccount() { 84  Account account = new Account(new ECKey()); 85  saveAccount(account); 86  return account.getAddress(); 87  } 88  89  public RskAddress addAccount(String passphrase) { 90  Account account = new Account(new ECKey()); 91  saveAccount(account, passphrase); 92  return account.getAddress(); 93  } 94  95  public RskAddress addAccount(Account account) { 96  saveAccount(account); 97  return account.getAddress(); 98  } 99  100  public Account getAccount(RskAddress addr) { 101  synchronized (accessLock) { 102  if (!accounts.containsKey(addr)) { 103  return null; 104  } 105  106  if (unlocksTimeouts.containsKey(addr)) { 107  long ending = unlocksTimeouts.get(addr); 108  long time = System.currentTimeMillis(); 109  if (ending < time) { 110  unlocksTimeouts.remove(addr); 111  accounts.remove(addr); 112  return null; 113  } 114  } 115  return new Account(ECKey.fromPrivate(accounts.get(addr))); 116  } 117  } 118  119  public Account getAccount(RskAddress addr, String passphrase) { 120  synchronized (accessLock) { 121  byte[] encrypted = keyDS.get(addr.getBytes()); 122  123  if (encrypted == null) { 124  return null; 125  } 126  127  return new Account(ECKey.fromPrivate(decryptAES(encrypted, passphrase.getBytes(StandardCharsets.UTF_8)))); 128  } 129  } 130  131  public boolean unlockAccount(RskAddress addr, String passphrase, long duration) { 132  long ending = System.currentTimeMillis() + duration; 133  boolean unlocked = unlockAccount(addr, passphrase); 134  135  if (unlocked) { 136  synchronized (accessLock) { 137  unlocksTimeouts.put(addr, ending); 138  } 139  } 140  141  return unlocked; 142  } 143  144  public boolean unlockAccount(RskAddress addr, String passphrase) { 145  Account account; 146  147  synchronized (accessLock) { 148  byte[] encrypted = keyDS.get(addr.getBytes()); 149  150  if (encrypted == null) { 151  return false; 152  } 153  154  account = new Account(ECKey.fromPrivate(decryptAES(encrypted, passphrase.getBytes(StandardCharsets.UTF_8)))); 155  } 156  157  saveAccount(account); 158  159  return true; 160  } 161  162  public boolean lockAccount(RskAddress addr) { 163  synchronized (accessLock) { 164  if (!accounts.containsKey(addr)) { 165  return false; 166  } 167  168  accounts.remove(addr); 169  return true; 170  } 171  } 172  173  public byte[] addAccountWithSeed(String seed) { 174  return addAccountWithPrivateKey(Keccak256Helper.keccak256(seed.getBytes(StandardCharsets.UTF_8))); 175  } 176  177  public byte[] addAccountWithPrivateKey(byte[] privateKeyBytes) { 178  Account account = new Account(ECKey.fromPrivate(privateKeyBytes)); 179  synchronized (accessLock) { 180  RskAddress addr = addAccount(account); 181  if (!this.initialAccounts.contains(addr)) { 182  this.initialAccounts.add(addr); 183  } 184  return addr.getBytes(); 185  } 186  } 187  188  public RskAddress addAccountWithPrivateKey(byte[] privateKeyBytes, String passphrase) { 189  Account account = new Account(ECKey.fromPrivate(privateKeyBytes)); 190  191  saveAccount(account, passphrase); 192  193  return account.getAddress(); 194  } 195  196  private void saveAccount(Account account) { 197  synchronized (accessLock) { 198  accounts.put(account.getAddress(), account.getEcKey().getPrivKeyBytes()); 199  } 200  } 201  202  private void saveAccount(Account account, String passphrase) { 203  byte[] address = account.getAddress().getBytes(); 204  byte[] privateKeyBytes = account.getEcKey().getPrivKeyBytes(); 205  byte[] encrypted = encryptAES(privateKeyBytes, passphrase.getBytes(StandardCharsets.UTF_8)); 206  207  synchronized (accessLock) { 208  keyDS.put(address, encrypted); 209  } 210  } 211  212  private byte[] decryptAES(byte[] encryptedBytes, byte[] passphrase) { 213  try { 214  ByteArrayInputStream in = new ByteArrayInputStream(encryptedBytes); 215  ObjectInputStream byteStream = new ObjectInputStream(in); 216  KeyCrypterAes keyCrypter = new KeyCrypterAes(); 217  KeyParameter keyParameter = new KeyParameter(Sha256Hash.hash(passphrase)); 218  219  ArrayList<byte[]> bytes = (ArrayList<byte[]>) byteStream.readObject(); 220  EncryptedData data = new EncryptedData(bytes.get(1), bytes.get(0)); 221  222  return keyCrypter.decrypt(data, keyParameter); 223  } catch (IOException | ClassNotFoundException e) { 224  //There are lines of code that should never be executed, this is one of those 225  throw new IllegalStateException(e); 226  } 227  } 228  229  private byte[] encryptAES(byte[] privateKeyBytes, byte[] passphrase) { 230  KeyCrypterAes keyCrypter = new KeyCrypterAes(); 231  KeyParameter keyParameter = new KeyParameter(Sha256Hash.hash(passphrase)); 232  EncryptedData enc = keyCrypter.encrypt(privateKeyBytes, keyParameter); 233  234  try { 235  ByteArrayOutputStream encryptedResult = new ByteArrayOutputStream(); 236  ObjectOutputStream byteStream = new ObjectOutputStream(encryptedResult); 237  238  ArrayList<byte[]> bytes = new ArrayList<>(); 239  bytes.add(enc.encryptedBytes); 240  bytes.add(enc.initialisationVector); 241  byteStream.writeObject(bytes); 242  243  return encryptedResult.toByteArray(); 244  } catch (IOException e) { 245  //How is this even possible ??? 246  throw new IllegalStateException(e); 247  } 248  } 249 }