1import stem 2import time 3import os 4import shutil 5 6from stem.response import ControlMessage 7 8from vanguards.control import get_consensus_weights 9 10import vanguards.vanguards 11from vanguards.vanguards import VanguardState 12from vanguards.vanguards import ExcludeNodes 13from vanguards.vanguards import _SEC_PER_HOUR 14 15from vanguards.vanguards import NUM_LAYER3_GUARDS 16from vanguards.vanguards import NUM_LAYER2_GUARDS 17from vanguards.vanguards import MIN_LAYER3_LIFETIME_HOURS 18from vanguards.vanguards import MAX_LAYER3_LIFETIME_HOURS 19from vanguards.vanguards import MIN_LAYER2_LIFETIME_HOURS 20from vanguards.vanguards import MAX_LAYER2_LIFETIME_HOURS 21 22try: 23 xrange 24except NameError: 25 xrange = range 26 27def replacement_checks(state, routers, weights): 28 remove2_idhex = state.layer2[0].idhex 29 remove3_idhex = state.layer3[0].idhex 30 31 # - Remove a layer2 guard from it 32 # - Remove a layer3 guard from it 33 routers = list(filter(lambda x: x.fingerprint != remove2_idhex \ 34 and x.fingerprint != remove3_idhex, 35 routers)) 36 37 assert remove2_idhex in map(lambda x: x.idhex, state.layer2) 38 assert remove3_idhex in map(lambda x: x.idhex, state.layer3) 39 keep2 = map(lambda x: x.idhex, 40 filter(lambda x: x.idhex != remove2_idhex \ 41 and x.idhex != remove3_idhex, 42 state.layer2)) 43 keep3 = map(lambda x: x.idhex, 44 filter(lambda x: x.idhex != remove2_idhex \ 45 and x.idhex != remove3_idhex, 46 state.layer3)) 47 state.consensus_update(routers, weights, ExcludeNodes(MockController())) 48 sanity_check(state) 49 assert not remove2_idhex in map(lambda x: x.idhex, state.layer2) 50 assert not remove3_idhex in map(lambda x: x.idhex, state.layer3) 51 for k in keep2: assert k in map(lambda x: x.idhex, state.layer2) 52 for k in keep3: assert k in map(lambda x: x.idhex, state.layer3) 53 54 remove2_idhex = state.layer2[1].idhex 55 remove3_idhex = state.layer3[1].idhex 56 57 # - Mark a layer2 guard way in the past 58 # - Mark a layer3 guard way in the past 59 state.layer2[1].expires_at = time.time() - 10 60 state.layer3[1].expires_at = time.time() - 10 61 62 assert remove2_idhex in map(lambda x: x.idhex, state.layer2) 63 assert remove3_idhex in map(lambda x: x.idhex, state.layer3) 64 keep2 = map(lambda x: x.idhex, 65 filter(lambda x: x.idhex != remove2_idhex, 66 state.layer2)) 67 keep3 = map(lambda x: x.idhex, 68 filter(lambda x: x.idhex != remove3_idhex, 69 state.layer3)) 70 state.consensus_update(routers, weights, ExcludeNodes(MockController())) 71 sanity_check(state) 72 assert not remove2_idhex in map(lambda x: x.idhex, state.layer2) 73 assert not remove3_idhex in map(lambda x: x.idhex, state.layer3) 74 for k in keep2: assert k in map(lambda x: x.idhex, state.layer2) 75 for k in keep3: assert k in map(lambda x: x.idhex, state.layer3) 76 77 # - Mark all guards way in the past 78 for g in state.layer2: 79 g.expires_at = time.time() - 10 80 for g in state.layer3: 81 g.expires_at = time.time() - 10 82 83 state.consensus_update(routers, weights, ExcludeNodes(MockController())) 84 sanity_check(state) 85 86 # Remove a node by idhex a few different ways 87 controller = MockController() 88 controller.exclude_nodes = \ 89 str(state.layer2[0].idhex)+","+str("$"+state.layer3[0].idhex)+","+\ 90 str(state.layer2[1].idhex+"~lol")+","+\ 91 str("$"+state.layer3[1].idhex+"~lol")+","+\ 92 str(state.layer2[2].idhex+"=lol")+","+\ 93 str("$"+state.layer3[2].idhex+"=lol") 94 95 removed2 = \ 96 [state.layer2[0].idhex, state.layer2[1].idhex, state.layer2[2].idhex] 97 removed3 = \ 98 [state.layer3[0].idhex, state.layer3[1].idhex, state.layer3[2].idhex] 99 100 for r in removed2: 101 assert r in map(lambda x: x.idhex, state.layer2) 102 for r in removed3: 103 assert r in map(lambda x: x.idhex, state.layer3) 104 105 keep3 = state.layer3[3].idhex 106 state.consensus_update(routers, weights, ExcludeNodes(controller)) 107 for r in removed2: 108 assert not r in map(lambda x: x.idhex, state.layer2) 109 for r in removed3: 110 assert not r in map(lambda x: x.idhex, state.layer3) 111 assert keep3 in map(lambda x: x.idhex, state.layer3) 112 113def sanity_check(state): 114 assert len(state.layer2) == NUM_LAYER2_GUARDS 115 assert len(state.layer3) == NUM_LAYER3_GUARDS 116 117 for g in state.layer2: 118 assert g.expires_at - g.chosen_at < MAX_LAYER2_LIFETIME_HOURS*_SEC_PER_HOUR 119 assert g.expires_at - g.chosen_at >= MIN_LAYER2_LIFETIME_HOURS*_SEC_PER_HOUR 120 121 for g in state.layer3: 122 assert g.expires_at - g.chosen_at < MAX_LAYER3_LIFETIME_HOURS*_SEC_PER_HOUR 123 assert g.expires_at - g.chosen_at >= MIN_LAYER3_LIFETIME_HOURS*_SEC_PER_HOUR 124 125class MockController: 126 def __init__(self): 127 self.exclude_nodes = None 128 self.exclude_unknown = "1" 129 self.got_set_conf = False 130 self.got_save_conf = False 131 self.get_info_vals = {} 132 133 # FIXME: os.path.join 134 def get_network_statuses(self): 135 return list(stem.descriptor.parse_file("tests/cached-microdesc-consensus", 136 document_handler = 137 stem.descriptor.DocumentHandler.ENTRIES)) 138 139 def get_conf(self, key): 140 if key == "DataDirectory": 141 return "tests" 142 if key == "ExcludeNodes": 143 return self.exclude_nodes 144 if key == "GeoIPExcludeUnknown": 145 return self.exclude_unknown 146 147 def set_conf(self, key, val): 148 self.got_set_conf = True 149 if key == "NumPrimaryGuards": 150 raise stem.InvalidArguments() 151 152 def save_conf(self): 153 self.got_save_conf = True 154 raise stem.OperationFailed("Bad") 155 156 def get_info(self, key, default=None): 157 if key in self.get_info_vals: 158 return self.get_info_vals[key] 159 else: 160 return default 161 162def test_new_vanguards(): 163 state = VanguardState("tests/state.mock2") 164 165 # - Load a routerlist using stem 166 routers = list(stem.descriptor.parse_file("tests/cached-microdesc-consensus", 167 document_handler = 168 stem.descriptor.DocumentHandler.ENTRIES)) 169 weights = get_consensus_weights("tests/cached-microdesc-consensus") 170 171 # - Perform basic rank checks from sort_and_index 172 (sorted_r, dict_r) = state.sort_and_index_routers(routers) 173 for i in xrange(len(sorted_r)-1): 174 assert sorted_r[i].measured >= sorted_r[i+1].measured 175 176 state.consensus_update(routers, weights, ExcludeNodes(MockController())) 177 sanity_check(state) 178 179 replacement_checks(state, routers, weights) 180 181def test_update_vanguards(): 182 controller = MockController() 183 vanguards.vanguards.LAYER1_LIFETIME_DAYS = 30 184 shutil.copy("tests/state.mock", "tests/state.mock.test") 185 state = VanguardState.read_from_file("tests/state.mock.test") 186 state.enable_vanguards = True 187 sanity_check(state) 188 189 state.new_consensus_event(controller, None) 190 sanity_check(state) 191 os.remove("tests/state.mock.test") 192 193 # test signal HUP 194 state.signal_event(controller, 195 ControlMessage.from_str("650 SIGNAL RELOAD\r\n", 196 "EVENT")) 197 sanity_check(state) 198 199def test_excludenodes(): 200 controller = MockController() 201 state = VanguardState("tests/state.mock2") 202 203 # - Load a routerlist using stem 204 routers = list(stem.descriptor.parse_file("tests/cached-microdesc-consensus", 205 document_handler = 206 stem.descriptor.DocumentHandler.ENTRIES)) 207 weights = get_consensus_weights("tests/cached-microdesc-consensus") 208 (sorted_r, dict_r) = state.sort_and_index_routers(routers) 209 210 state.consensus_update(routers, weights, ExcludeNodes(controller)) 211 sanity_check(state) 212 213 # * IP, CIDR, quad-mask 214 controller.exclude_nodes = \ 215 str(dict_r[state.layer2[0].idhex].address)+","+\ 216 str(dict_r[state.layer2[1].idhex].address)+"/24,"+\ 217 str(dict_r[state.layer2[2].idhex].address)+"/255.255.255.0" 218 removed2 = [state.layer2[0].idhex, state.layer2[1].idhex, 219 state.layer2[2].idhex] 220 221 for r in removed2: 222 assert r in map(lambda x: x.idhex, state.layer2) 223 state.consensus_update(routers, weights, ExcludeNodes(controller)) 224 sanity_check(state) 225 for r in removed2: 226 assert not r in map(lambda x: x.idhex, state.layer2) 227 228 # * GeoIP case mismatch 229 controller.exclude_nodes = "{Us}" 230 controller.exclude_unknown = "auto" 231 controller.get_info_vals["ip-to-country/"+dict_r[state.layer2[1].idhex].address] = "us" 232 controller.get_info_vals["ip-to-country/ipv4-available"] = "1" 233 removed2 = state.layer2[1].idhex 234 keep2 = state.layer2[0].idhex 235 state.consensus_update(routers, weights, ExcludeNodes(controller)) 236 237 sanity_check(state) 238 assert keep2 in map(lambda x: x.idhex, state.layer2) 239 assert not removed2 in map(lambda x: x.idhex, state.layer2) 240 241 # * Nicks 242 controller.exclude_nodes = \ 243 str(dict_r[state.layer2[0].idhex].nickname) 244 245 removed2 = state.layer2[0].idhex 246 keep2 = state.layer2[1].idhex 247 state.consensus_update(routers, weights, ExcludeNodes(controller)) 248 sanity_check(state) 249 assert not removed2 in map(lambda x: x.idhex, state.layer2) 250 assert keep2 in map(lambda x: x.idhex, state.layer2) 251 252 # FIXME: IPv6. Stem before 1.7.0 does not support IPv6 relays.. 253 254def test_disable(): 255 controller = MockController() 256 vanguards.vanguards.LAYER1_LIFETIME_DAYS = 30 257 shutil.copy("tests/state.mock", "tests/state.mock.test") 258 state = VanguardState.read_from_file("tests/state.mock.test") 259 state.enable_vanguards = False 260 sanity_check(state) 261 262 state.new_consensus_event(controller, None) 263 sanity_check(state) 264 assert controller.got_set_conf == False 265 assert controller.got_save_conf == False 266 os.remove("tests/state.mock.test") 267 268