1 package com.android.server.wifi; 2 3 import static org.mockito.Matchers.anyInt; 4 import static org.mockito.Matchers.anyString; 5 import static org.mockito.Mockito.mock; 6 import static org.mockito.Mockito.when; 7 8 import android.security.KeyStore; 9 import android.util.SparseArray; 10 11 import org.mockito.Matchers; 12 import org.mockito.invocation.InvocationOnMock; 13 import org.mockito.stubbing.Answer; 14 15 import java.util.Arrays; 16 import java.util.HashMap; 17 18 class MockKeyStore { 19 20 public static class KeyBlob { 21 public byte[] blob; 22 public int flag; 23 24 public void update(byte[] blob, int flag) { 25 this.blob = Arrays.copyOf(blob, blob.length); 26 this.flag = flag; 27 } 28 } 29 private SparseArray<HashMap<String, KeyBlob>> mStore; 30 31 public MockKeyStore() { 32 mStore = new SparseArray<HashMap<String, KeyBlob>>(); 33 } 34 35 public KeyStore createMock() { 36 KeyStore mock = mock(KeyStore.class); 37 when(mock.state()).thenReturn(KeyStore.State.UNLOCKED); 38 39 when(mock.put(anyString(), Matchers.any(byte[].class), anyInt(), anyInt())) 40 .thenAnswer(new Answer<Boolean>() { 41 42 @Override 43 public Boolean answer(InvocationOnMock invocation) throws Throwable { 44 Object[] args = invocation.getArguments(); 45 return put((String) args[0], (byte[]) args[1], (Integer) args[2], 46 (Integer) args[3]); 47 } 48 }); 49 50 when(mock.importKey(anyString(), Matchers.any(byte[].class), anyInt(), anyInt())) 51 .thenAnswer(new Answer<Boolean>() { 52 53 @Override 54 public Boolean answer(InvocationOnMock invocation) throws Throwable { 55 Object[] args = invocation.getArguments(); 56 return importKey((String) args[0], (byte[]) args[1], (Integer) args[2], 57 (Integer) args[3]); 58 } 59 }); 60 61 when(mock.delete(anyString(), anyInt())).thenAnswer(new Answer<Boolean>() { 62 63 @Override 64 public Boolean answer(InvocationOnMock invocation) throws Throwable { 65 Object[] args = invocation.getArguments(); 66 return delete((String) args[0], (Integer) args[1]); 67 } 68 }); 69 70 when(mock.contains(anyString(), anyInt())).thenAnswer(new Answer<Boolean>() { 71 72 @Override 73 public Boolean answer(InvocationOnMock invocation) throws Throwable { 74 Object[] args = invocation.getArguments(); 75 return contains((String) args[0], (Integer) args[1]); 76 } 77 }); 78 79 return mock; 80 } 81 82 private KeyBlob access(int uid, String key, boolean createIfNotExist) { 83 if (mStore.get(uid) == null) { 84 mStore.put(uid, new HashMap<String, KeyBlob>()); 85 } 86 HashMap<String, KeyBlob> map = mStore.get(uid); 87 if (map.containsKey(key)) { 88 return map.get(key); 89 } else { 90 if (createIfNotExist) { 91 KeyBlob blob = new KeyBlob(); 92 map.put(key, blob); 93 return blob; 94 } else { 95 return null; 96 } 97 } 98 } 99 100 public KeyBlob getKeyBlob(int uid, String key) { 101 return access(uid, key, false); 102 } 103 104 private boolean put(String key, byte[] value, int uid, int flags) { 105 access(uid, key, true).update(value, flags); 106 return true; 107 } 108 109 private boolean importKey(String keyName, byte[] key, int uid, int flags) { 110 return put(keyName, key, uid, flags); 111 } 112 113 private boolean delete(String key, int uid) { 114 if (mStore.get(uid) != null) { 115 mStore.get(uid).remove(key); 116 } 117 return true; 118 } 119 120 private boolean contains(String key, int uid) { 121 return access(uid, key, false) != null; 122 } 123 124 @Override 125 public String toString() { 126 StringBuilder sb = new StringBuilder(); 127 sb.append("KeyStore {"); 128 for (int i = 0; i < mStore.size(); i++) { 129 int uid = mStore.keyAt(i); 130 for (String keyName : mStore.get(uid).keySet()) { 131 sb.append(String.format("%d:%s, ", uid, keyName)); 132 } 133 } 134 sb.append("}"); 135 return sb.toString(); 136 } 137 } 138
