|
| 1 | +import time |
| 2 | + |
| 3 | +from slips_files.common.slips_utils import utils |
| 4 | +from slips_files.common.style import green |
| 5 | + |
| 6 | + |
| 7 | +class DoSProtector: |
| 8 | + def __init__(self, input): |
| 9 | + self.input = input |
| 10 | + self.db = self.input.db |
| 11 | + self.is_running_non_stop: bool = self.db.is_running_non_stop() |
| 12 | + # is slips (input.py) is given > this number of flow per min, |
| 13 | + # this protector runs |
| 14 | + self.flows_per_min_threshold = 2000 |
| 15 | + self.flow_sampling_stop_time = 0 |
| 16 | + # number of seconds slips is going to be skipping flows for before |
| 17 | + # returning to normal (aka before going back to reading all flows) |
| 18 | + self.sampling_time_window = 60 |
| 19 | + self._is_now_sampling = False |
| 20 | + |
| 21 | + def _get_input_flows_per_min(self) -> int: |
| 22 | + input_flows_per_s = ( |
| 23 | + self.db.get_core_module_flows_per_second("Input") or 0 |
| 24 | + ) |
| 25 | + input_flows_per_min = int(input_flows_per_s) * 60 |
| 26 | + return input_flows_per_min |
| 27 | + |
| 28 | + def _get_sampling_ratio(self) -> int: |
| 29 | + """ |
| 30 | + sr = flow_per_min² / 20000 |
| 31 | + this sr is the number of flows we're gonna skip to protect slips |
| 32 | + from DoS (or high traffic in general) |
| 33 | + """ |
| 34 | + input_flows_per_min = self._get_input_flows_per_min() |
| 35 | + if not input_flows_per_min: |
| 36 | + return 1 |
| 37 | + |
| 38 | + return input_flows_per_min**2 / 20000 |
| 39 | + |
| 40 | + def _should_run(self) -> bool: |
| 41 | + """ |
| 42 | + Returns true if slips is under high traffic and the DoS protector |
| 43 | + should run. |
| 44 | + Runs only when analysing an interface or a growing zeek dir. |
| 45 | + return True if: |
| 46 | + 1. if high traffic is detected |
| 47 | + 2. we're in the 1 min window after slips has detected a high |
| 48 | + traffic. this is the 1 min of skipping flows before rechecking if |
| 49 | + the read number of flows has decreased. |
| 50 | + """ |
| 51 | + if not self.is_running_non_stop: |
| 52 | + return False |
| 53 | + |
| 54 | + if time.time() < self.flow_sampling_stop_time: |
| 55 | + # we should still be sampling. |
| 56 | + return True |
| 57 | + |
| 58 | + input_flows_per_min = self._get_input_flows_per_min() |
| 59 | + should_skip_flows = input_flows_per_min > self.flows_per_min_threshold |
| 60 | + |
| 61 | + if self._is_now_sampling and input_flows_per_min == 0: |
| 62 | + # this means we justtt stopped sampling, now we want slips to |
| 63 | + # keep thinking thta it's in a sampling state until we get a |
| 64 | + # input_flows_per_min = something, once we have a number we can |
| 65 | + # decide whether to stop sampling or not, but until then we want to keep the sampling state |
| 66 | + pass |
| 67 | + elif ( |
| 68 | + not should_skip_flows |
| 69 | + and self._is_now_sampling |
| 70 | + and input_flows_per_min |
| 71 | + ): |
| 72 | + # slips was sampling and now stopped officially stopped, |
| 73 | + # we have a input_flows_per_min that's less than the threshold. |
| 74 | + self._is_now_sampling = False |
| 75 | + self.input.print( |
| 76 | + f"Throughput is back to normal. Input " |
| 77 | + f"flows/min = {green(input_flows_per_min)}. " |
| 78 | + f"Slips stopped skipping flows." |
| 79 | + ) |
| 80 | + |
| 81 | + return should_skip_flows |
| 82 | + |
| 83 | + def _update_flow_sampling_stop_time_if_needed(self) -> bool: |
| 84 | + """ |
| 85 | + sets the next stop time to |
| 86 | + now + sampling_time_window |
| 87 | + if the time now exceeded the last registered flow_sampling_stop_time |
| 88 | + """ |
| 89 | + if time.time() > self.flow_sampling_stop_time: |
| 90 | + # flow sampling is going to take place for the next 1 min |
| 91 | + self.flow_sampling_stop_time = ( |
| 92 | + time.time() + self.sampling_time_window |
| 93 | + ) |
| 94 | + return True |
| 95 | + return False |
| 96 | + |
| 97 | + def get_number_of_flows_to_skip_and_time_to_stop_sampling(self) -> int: |
| 98 | + if not self._should_run(): |
| 99 | + return 0 |
| 100 | + |
| 101 | + sampling_time_updated = ( |
| 102 | + self._update_flow_sampling_stop_time_if_needed() |
| 103 | + ) |
| 104 | + |
| 105 | + # -1 means read 1 flow every sampling_ratio flows. |
| 106 | + # at 2000 flows/min → sr = 200, read 1 flow every 200 flows |
| 107 | + # at 3000 flows/min → sr = 450, read 1 flow every 450 flows |
| 108 | + # at 4000 flows/min → sr = 800, etc. |
| 109 | + sampling_ratio = int(self._get_sampling_ratio() - 1) |
| 110 | + self.print_skipping_flows_warning( |
| 111 | + sampling_ratio, sampling_time_updated |
| 112 | + ) |
| 113 | + |
| 114 | + return sampling_ratio |
| 115 | + |
| 116 | + def print_skipping_flows_warning( |
| 117 | + self, sampling_ratio: int, sampling_time_updated: bool |
| 118 | + ): |
| 119 | + """Prints a warning every time slips decides to start sampling |
| 120 | + again""" |
| 121 | + if sampling_time_updated and sampling_ratio: |
| 122 | + sr = green(f"1/{sampling_ratio}") |
| 123 | + human_readable_time_to_stop_sampling = utils.convert_ts_format( |
| 124 | + self.flow_sampling_stop_time, utils.alerts_format |
| 125 | + ) |
| 126 | + green_time_to_stop_sampling = green( |
| 127 | + human_readable_time_to_stop_sampling |
| 128 | + ) |
| 129 | + if self._is_now_sampling: |
| 130 | + # slips decided to extend the sampling period |
| 131 | + self.input.print( |
| 132 | + f"Slips is still under high " |
| 133 | + f"traffic. The time to stop sampling has been extended to " |
| 134 | + f"{green_time_to_stop_sampling} " |
| 135 | + ) |
| 136 | + else: |
| 137 | + # reaching here means slips decided again to start sampling flows |
| 138 | + self.input.print( |
| 139 | + f"Slips started skipping flows due to high " |
| 140 | + f"traffic for DoS protection. " |
| 141 | + f"Sampling ratio: {sr} flows. " |
| 142 | + f"Time to stop sampling: {green_time_to_stop_sampling} " |
| 143 | + ) |
| 144 | + self._is_now_sampling = True |
0 commit comments