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