phonerr/
lib.rs

1pub mod util;
2
3use std::{cmp::Ordering, fmt::Debug, os::unix::fs::MetadataExt, path::PathBuf};
4
5use anyhow::anyhow;
6use colored::Colorize;
7use tokio::{
8    fs::{self, File},
9    io::{AsyncReadExt, AsyncSeekExt},
10};
11use util::{download_file, download_file_silent, replace_home, u8_i32};
12
13#[derive(Debug)]
14pub struct PhoneData {
15    pub data_path: PathBuf,
16    pub version: String,
17    pub start_index: u32,
18    pub record: Vec<u8>,
19    pub index: Vec<u8>,
20    pub size: u64,
21}
22
23// 索引数据
24#[derive(Debug)]
25pub struct IndexData {
26    pub phone_prefix: i32,
27    pub offset: i32,
28    pub isp: u8,
29}
30
31#[derive(Debug)]
32pub struct RecordData {
33    pub province: String,
34    pub city: String,
35    pub zip_code: String,
36    pub area_code: String,
37    pub isp: String,
38}
39
40impl RecordData {
41    pub fn display(&self) {
42        if self.province != self.city {
43            println!(
44                "{} {}{} {}",
45                format!("[{}]", self.area_code).magenta(),
46                self.province.green(),
47                self.city.green(),
48                self.isp.red().bold()
49            );
50        } else {
51            println!(
52                "{} {} {}",
53                format!("[{}]", self.area_code).magenta(),
54                self.province.green(),
55                self.isp.red().bold()
56            );
57        }
58    }
59}
60
61// 运营商
62enum Isp {
63    Cmcc = 1,
64    Cucc = 2,
65    Ctcc = 3,
66    CtccV = 4,
67    CuccV = 5,
68    CmccV = 6,
69    Cbcc = 7,
70    CbccV = 8,
71}
72
73impl Isp {
74    fn from_u8(i: u8) -> Option<Isp> {
75        match i {
76            1 => Some(Isp::Cmcc),
77            2 => Some(Isp::Cucc),
78            3 => Some(Isp::Ctcc),
79            4 => Some(Isp::CtccV),
80            5 => Some(Isp::CuccV),
81            6 => Some(Isp::CmccV),
82            7 => Some(Isp::Cbcc),
83            8 => Some(Isp::CbccV),
84            _ => None,
85        }
86    }
87
88    fn get_name(&self) -> String {
89        match self {
90            Isp::Cmcc => "中国移动".to_string(),
91            Isp::Cucc => "中国联通".to_string(),
92            Isp::Ctcc => "中国电信".to_string(),
93            Isp::CtccV => "中国电信虚拟运营商".to_string(),
94            Isp::CuccV => "中国联通虚拟运营商".to_string(),
95            Isp::CmccV => "中国移动虚拟运营商".to_string(),
96            Isp::Cbcc => "中国广电".to_string(),
97            Isp::CbccV => "中国广电虚拟运营商".to_string(),
98        }
99    }
100}
101
102const PHONE_DATA_URL: &str = "https://raw.githubusercontent.com/ls0f/phone/master/phone/phone.dat";
103
104impl PhoneData {
105    pub fn new(data_path: Option<&str>) -> Self {
106        let data_path = PathBuf::from(replace_home(
107            data_path.unwrap_or(Self::get_data_path().as_str()),
108        ));
109        Self {
110            data_path,
111            version: String::new(),
112            start_index: 0,
113            record: vec![],
114            index: vec![],
115            size: 0,
116        }
117    }
118
119    pub fn get_data_path() -> String {
120        replace_home("~/.cache/phoner/phone.dat")
121    }
122
123    pub async fn download_file(
124        &self,
125        download_url: Option<String>,
126        silent: bool,
127    ) -> Result<(), anyhow::Error> {
128        if !self.data_path.parent().unwrap().exists() {
129            fs::create_dir_all(self.data_path.parent().unwrap()).await?;
130        }
131        if silent {
132            download_file_silent(
133                download_url.unwrap_or(PHONE_DATA_URL.to_string()).as_str(),
134                &PathBuf::from(Self::get_data_path()),
135            )
136            .await?;
137        } else {
138            download_file(
139                download_url.unwrap_or(PHONE_DATA_URL.to_string()).as_str(),
140                &PathBuf::from(Self::get_data_path()),
141            )
142            .await?;
143        }
144        Ok(())
145    }
146
147    async fn init(&mut self) -> Result<(), anyhow::Error> {
148        if !self.data_path.exists() {
149            self.download_file(None, true).await?;
150        }
151        let mut file = File::open(self.data_path.clone()).await?;
152        // 读取头部8个字节 版本号, 第一个索引的偏移
153        let mut version_bytes: [u8; 4] = [0; 4];
154        file.read_exact(&mut version_bytes).await?;
155        let version = String::from_utf8_lossy(&version_bytes);
156        self.version = version.to_string();
157        file.seek(std::io::SeekFrom::Start(4)).await?;
158        // 小端字节序
159        let start_index = file.read_u32_le().await?;
160        self.start_index = start_index;
161        // 读取记录区[8:start_index]
162        file.seek(std::io::SeekFrom::Start(8)).await?;
163        let mut record = vec![0u8; (start_index - 8) as usize];
164        file.read_exact(&mut record).await?;
165        self.record = record;
166        // 读取索引区
167        let mut index = vec![];
168        file.read_to_end(&mut index).await?;
169        self.index = index;
170        let metadata = file.metadata().await.unwrap();
171        self.size = metadata.size();
172        Ok(())
173    }
174
175    fn search_index(&self, phone: &str) -> Result<Option<IndexData>, anyhow::Error> {
176        if self.index.is_empty() {
177            return Err(anyhow!("数据未初始化"));
178        }
179        let phone_prefix: i32 = phone[..7].parse()?;
180        // 采用二分法查找index
181        let mut start = 0;
182        let mut end = self.index.len() / 9 - 1;
183        let mut position = None;
184        loop {
185            let mid = (start + end) / 2;
186            // 前4位为手机号
187            let prefix = &self.index[mid * 9..(mid * 9 + 4)];
188            let prefix = u8_i32(prefix);
189            match prefix.cmp(&phone_prefix) {
190                Ordering::Greater => {
191                    if mid == 0 {
192                        break;
193                    }
194                    end = mid - 1;
195                }
196                Ordering::Equal => {
197                    position = Some(mid);
198                    break;
199                }
200                Ordering::Less => {
201                    start = mid + 1;
202                }
203            }
204            if start > end {
205                break;
206            }
207        }
208        if let Some(pos) = position {
209            let data = &self.index[pos * 9..(pos + 1) * 9];
210            return Ok(Some(IndexData {
211                phone_prefix: u8_i32(&data[..4]),
212                offset: u8_i32(&data[4..8]),
213                isp: u8_i32(&data[8..]) as u8,
214            }));
215        }
216        Ok(None)
217    }
218
219    pub async fn query(&mut self, phone: &str, init: bool) -> Result<RecordData, anyhow::Error> {
220        if init {
221            self.init().await?;
222        }
223        let index = self.search_index(phone).unwrap();
224        if index.is_none() {
225            return Err(anyhow!("未找到数据"));
226        }
227        let index = index.unwrap();
228        let mut record = self.read_recrod(index.offset).unwrap();
229        if let Some(isp) = Isp::from_u8(index.isp) {
230            record.isp = isp.get_name();
231        }
232        Ok(record)
233    }
234
235    fn read_recrod(&self, offset: i32) -> Result<RecordData, anyhow::Error> {
236        let res: Vec<u8> = self
237            .record
238            .iter()
239            .skip((offset - 8) as usize)
240            .take_while(|x| **x != 0)
241            .copied()
242            .collect();
243        let data = String::from_utf8(res).unwrap();
244        let data_split: Vec<&str> = data.split('|').collect();
245        if data_split.len() != 4 {
246            return Err(anyhow!("数据格式错误: {}", data));
247        }
248        Ok(RecordData {
249            province: data_split[0].to_string(),
250            city: data_split[1].to_string(),
251            zip_code: data_split[2].to_string(),
252            area_code: data_split[3].to_string(),
253            isp: String::new(),
254        })
255    }
256
257    pub async fn print_db_info(&mut self) -> Result<(), anyhow::Error> {
258        if !self.data_path.exists() {
259            return Err(anyhow!("手机号码库不存在"));
260        }
261        self.init().await?;
262        println!("{:<7}: {}", "version".green().bold(), self.version,);
263        println!("{:<7}: {}", "lines".green().bold(), self.index.len() / 9);
264        println!(
265            "{:<7}: {} bytes ({:.4}MB)",
266            "size".green().bold(),
267            self.size,
268            self.size as f64 / 1024.0 / 1024.0,
269        );
270        Ok(())
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use crate::PhoneData;
277
278    #[tokio::test]
279    async fn test_query() -> Result<(), anyhow::Error> {
280        let mut phone_data = PhoneData::new(Some("./phone.dat"));
281        let record = phone_data.query("13456789012", true).await?;
282        assert_eq!(record.province, "浙江");
283        assert_eq!(record.city, "杭州");
284        assert_eq!(record.isp, "中国移动");
285        assert_eq!(record.zip_code, "310000");
286        assert_eq!(record.area_code, "0571");
287        Ok(())
288    }
289}