| 392 | |
| 393 | |
| 394 | class Rule (object): |
| 395 | def __init__(self, rule_type, domain, ips, rebinds, threshold): |
| 396 | self.type = rule_type |
| 397 | self.domain = domain |
| 398 | self.ips = ips |
| 399 | self.rebinds = rebinds |
| 400 | self.rebind_threshold = threshold |
| 401 | |
| 402 | # we need an additional object to track the rebind rules |
| 403 | if self.rebinds is not None: |
| 404 | self.match_history = {} |
| 405 | self.rebinds = self._round_robin(rebinds) |
| 406 | self.ips = self._round_robin(ips) |
| 407 | |
| 408 | def _round_robin(self, ip_list): |
| 409 | """ |
| 410 | Creates a generator over a list modulo list length to equally move between all elements in the list each request |
| 411 | Since we have rules broken out into objects now, we can have this without much overhead. |
| 412 | """ |
| 413 | # check to make sure we don't try to modulo by zero |
| 414 | # if we would, just add the same element to the list again. |
| 415 | if len(ip_list) == 1: |
| 416 | ip_list.append(ip_list[0]) |
| 417 | |
| 418 | # should be fine to continue now. |
| 419 | index = 0 |
| 420 | while 1: # never stop iterating - it's OK since we dont always run |
| 421 | yield ip_list[index] |
| 422 | index += 1 |
| 423 | index = index % len(ip_list) |
| 424 | |
| 425 | def match(self, req_type, domain, addr): |
| 426 | # assert that the query type and domain match |
| 427 | try: |
| 428 | req_type = TYPE[req_type] |
| 429 | except KeyError: |
| 430 | return None |
| 431 | |
| 432 | try: |
| 433 | assert self.type == req_type |
| 434 | except AssertionError: |
| 435 | return None |
| 436 | |
| 437 | try: |
| 438 | assert self.domain.match(domain.decode()) |
| 439 | except AssertionError: |
| 440 | return None |
| 441 | |
| 442 | # Check to see if we have a rebind rule and if we do, return that addr first |
| 443 | if self.rebinds: |
| 444 | if self.match_history.get(addr) is not None: |
| 445 | |
| 446 | # passed the threshold - start doing a rebind |
| 447 | if self.match_history[addr] >= self.rebind_threshold: |
| 448 | return next(self.rebinds) |
| 449 | |
| 450 | # plus one |
| 451 | else: |