1pub mod ip_lookup {
2
3 use std::net::Ipv4Addr;
4 use std::path::PathBuf;
5 use std::str::FromStr;
6 use std::cmp::Ordering;
7 use std::error::Error;
8 use csv::Reader;
9
10
11 #[derive(Debug, Clone)]
12 pub struct IpRange {
13 start: u32,
14 end: u32,
15 pub country: String,
16 pub region: String,
17 pub city: String,
18 }
19
20 #[derive(Debug)]
21 pub struct Looker {
22 pub file_path: PathBuf,
23 pub ip_ranges: Vec<IpRange>,
24 }
25
26 #[derive(Debug, Default)]
27 pub struct LookerBuilder {
28 file_path: Option<PathBuf>,
29 allowed_countries: Option<Vec<String>>,
30 }
31
32 pub trait IpLookup {
33 fn look_up(&self, ip: &str) -> Option<IpRange>;
34 fn look_up_ipv4(&self, ip: &Ipv4Addr) -> Option<IpRange>;
35 }
36
37 impl Looker {
38
39 pub fn new(file_path: PathBuf) -> Self {
40
41 let ip_ranges = match read_ip_ranges(file_path.to_str().expect("IP CSV file not found"), None) {
42 Ok(ranges) => ranges,
43 Err(e) => {
44 log::error!("Error reading IP ranges: {}", e);
45 Vec::new()
46 }
47 };
48 Looker {
49 file_path,
50 ip_ranges,
51 }
52
53 }
54
55 pub fn builder() -> LookerBuilder {
56 LookerBuilder::new()
57 }
58
59 }
60
61 impl LookerBuilder {
62 pub fn new() -> Self {
63 Self::default()
64 }
65
66 pub fn file_path(mut self, path: PathBuf) -> Self {
67 self.file_path = Some(path);
68 self
69 }
70
71 pub fn allowed_countries(mut self, countries: Vec<String>) -> Self {
72 if countries.is_empty() {
73 log::warn!("Allowed countries is empty, filter will be ignored!");
74 self.allowed_countries = None;
75 return self;
76 }
77
78 self.allowed_countries = Some(countries);
79 self
80 }
81
82 pub fn build(self) -> Result<Looker, Box<dyn Error>> {
83 let ip_ranges = match read_ip_ranges(self.file_path.as_ref().expect("IP CSV file not found").to_str().expect("Invalid file path"), self.allowed_countries.as_ref()) {
84 Ok(ranges) => ranges,
85 Err(e) => {
86 log::error!("Error reading IP ranges: {}", e);
87 Vec::new()
88 }
89 };
90
91 Ok(Looker {
92 file_path: self.file_path.unwrap(),
93 ip_ranges,
94 })
95 }
96 }
97
98 fn read_ip_ranges(file_path: &str, allowed_countries: Option<&Vec<String>>) -> Result<Vec<IpRange>, Box<dyn Error>> {
99 let mut rdr = Reader::from_path(file_path)?;
100 let mut ip_ranges = Vec::new();
101
102 let allowed_countries = match allowed_countries {
103 Some(filter) => {
104 if filter.is_empty() {
105 log::warn!("Country filter is empty, filter will be ignored!");
106 None
107 } else {
108 Some(filter)
109 }
110 },
111 None => None,
112 };
113
114 for result in rdr.records() {
115 let record = result?;
116 let start: u32 = record[0].parse()?;
117 let end: u32 = record[1].parse()?;
118 let country = record[2].to_string();
119 let region = record[4].to_string();
120 let city = record[5].to_string();
121
122 if let Some(ref filter) = allowed_countries {
123 if !filter.contains(&country) {
124 continue;
125 }
126 }
127
128 ip_ranges.push(IpRange { start, end, country, region, city });
129 }
130
131 Ok(ip_ranges)
132 }
133
134 fn find_ip_range(ip: u32, ranges: &[IpRange]) -> Option<IpRange> {
135 ranges.binary_search_by(|range| {
136 if ip < range.start {
137 Ordering::Greater } else if ip > range.end {
139 Ordering::Less } else {
141 Ordering::Equal }
143 }).ok().map(|index| ranges[index].clone())
144 }
145
146 fn ip_string_to_decimal(ip: &str) -> Result<u32, String> {
147 let ip = Ipv4Addr::from_str(ip);
148 if ip.is_err() {
149 return Err("Invalid IP address".into());
150 }
151 let ip = ip.unwrap();
152 ip_to_decimal(&ip)
153 }
154
155 fn ip_to_decimal(ip: &Ipv4Addr) -> Result<u32,String> {
156 let octets = ip.octets();
157 let decimal = (octets[0] as u32) << 24
158 | (octets[1] as u32) << 16
159 | (octets[2] as u32) << 8
160 | octets[3] as u32;
161 Ok(decimal)
162 }
163
164 pub fn look_up(ip: &str, file_path: &str) -> Option<IpRange> {
165 let ip_decimal_to_use = match ip_string_to_decimal(ip) {
166 Err(e) => {
167 log::error!("Error: {}", e);
168 return None;
169 },
170 Ok(ip_decimal) => {
171 ip_decimal
172 }
173 };
174 let ip_ranges_to_use = match read_ip_ranges(file_path, None) {
175 Err(e) => {
176 log::error!("Error: {}", e);
177 return None;
178 },
179 Ok(ip_ranges) => {
180 ip_ranges
181 }
182 };
183
184 match find_ip_range(ip_decimal_to_use, &ip_ranges_to_use[..]) {
185 Some(range) => {
186 log::trace!("IP is in range: {:?}", range);
187 Some(range)
188 },
189 None => {
190 log::trace!("IP not found in any range");
191 None
192 }
193 }
194 }
195
196 pub fn look_up_filtered(ip: &str, file_path: &str, allowed_countries: &Vec<String>) -> Option<IpRange> {
197 let ip_decimal_to_use = match ip_string_to_decimal(ip) {
198 Err(e) => {
199 log::error!("Error: {}", e);
200 return None;
201 },
202 Ok(ip_decimal) => {
203 ip_decimal
204 }
205 };
206 let ip_ranges_to_use = match read_ip_ranges(file_path, Some(allowed_countries)) {
207 Err(e) => {
208 log::error!("Error: {}", e);
209 return None;
210 },
211 Ok(ip_ranges) => {
212 ip_ranges
213 }
214 };
215
216 match find_ip_range(ip_decimal_to_use, &ip_ranges_to_use[..]) {
217 Some(range) => {
218 log::trace!("IP is in range: {:?}", range);
219 Some(range)
220 },
221 None => {
222 log::trace!("IP not found in any range");
223 None
224 }
225 }
226 }
227
228 impl IpLookup for Looker {
229
230 fn look_up(&self, ip: &str) -> Option<IpRange> {
231 let ip = Ipv4Addr::from_str(ip);
232 match ip {
233 Err(e) => {
234 log::error!("Error: {}", e);
235 None
236 },
237 Ok(ip) => {
238 self.look_up_ipv4(&ip)
239 }
240 }
241
242 }
243
244 fn look_up_ipv4(&self, ip: &Ipv4Addr) -> Option<IpRange> {
245
246 let ip_decimal_to_use = match ip_to_decimal(ip) {
247 Err(e) => {
248 log::error!("Error: {}", e);
249 return None;
250 },
251 Ok(ip_decimal) => {
252 ip_decimal
253 }
254 };
255 let ip_ranges_to_use = &self.ip_ranges;
256
257 match find_ip_range(ip_decimal_to_use, &ip_ranges_to_use[..]) {
258 Some(range) => {
259 log::trace!("IP is in range: {:?}", range);
260 Some(range)
261 },
262 None => {
263 log::trace!("IP not found in any range");
264 None
265 }
266 }
267 }
268
269 }
270
271}
272
273pub use crate::ip_lookup::{look_up, look_up_filtered, Looker, LookerBuilder, IpLookup, };