1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

//! 基于 rust 实现的 zabbix proxy，实现了基本的代理功能。
//!
#[macro_use]
extern crate serde_derive;

#[macro_use]
extern crate log;

use failure::Error;

type Result<T> = std::result::Result<T, Error>;

mod request;
pub use self::request::{ZabbixDiscovery, ZabbixHost, ZabbixMetric, ZabbixRequest};

mod response;
pub use self::response::{Host, Item, ProxyResponse, Response};

use serde_json::Value;
use std::collections::{HashMap, HashSet};
use zabbix::ZabbixProtocol;

#[derive(Debug, Clone)]
pub struct ZabbixProxy {
    name: String,
    proto: ZabbixProtocol,
}

impl ZabbixProxy {
    pub const PROXY_CONFIG: &'static str = "proxy config";
    pub const HISTORY_DATA: &'static str = "history data";
    pub const PROXY_HEARTBEAT: &'static str = "proxy heartbeat";
    pub const AUTO_REGISTRATION: &'static str = "auto registration";

    pub fn new(name: &str, server: &str, port: u16) -> Self {
        let name = String::from(name);
        let proto = ZabbixProtocol::new(server, port);
        Self { name, proto }
    }

    fn send_request(&self, req: &ZabbixRequest, is_config: bool) -> Result<ProxyResponse> {
        let read_data = self.proto.send(&req.str())?;
        let response = if is_config {
            ProxyResponse::CONFIG(serde_json::from_slice(&read_data)?)
        } else {
            ProxyResponse::RESPONSE(serde_json::from_slice(&read_data)?)
        };

        Ok(response)
    }

    ///
    /// 从ZABBNIX服务端获取代理配置信息
    ///
    pub fn get_config(&self) -> Option<Value> {
        let req = ZabbixRequest::new(Self::PROXY_CONFIG, &self.name, Value::Null);
        if let Ok(r) = self.send_request(&req, true) {
            if let ProxyResponse::CONFIG(c) = r {
                return Some(c);
            }
        }
        None
    }

    ///
    /// 自动注册主机
    ///
    pub fn auto_register(&self, hosts: Vec<ZabbixHost>) -> Result<bool> {
        let hosts = serde_json::to_value(hosts)?;
        let req = ZabbixRequest::new(Self::AUTO_REGISTRATION, &self.name, hosts);
        if let Ok(r) = self.send_request(&req, false) {
            if let ProxyResponse::RESPONSE(c) = r {
                return Ok(c.success());
            }
        }
        Ok(false)
    }

    ///
    /// 向服务端发送心跳信息
    ///
    pub fn heart_beat(&self) -> Result<bool> {
        let req = ZabbixRequest::new(Self::PROXY_HEARTBEAT, &self.name, Value::Null);
        if let Ok(r) = self.send_request(&req, false) {
            if let ProxyResponse::RESPONSE(c) = r {
                return Ok(c.success());
            }
        }
        Ok(false)
    }

    ///
    /// 向服务端发送历史数据
    ///
    pub fn send_data(&self, data: &[ZabbixMetric]) -> Result<bool> {
        trace!("request = {:?}", Self::HISTORY_DATA);
        if !data.is_empty() {
            trace!("key[0] = {:?}", &data[0].key);
            trace!("host[0] = {:?}", &data[0].host);
            trace!("data[0] = {:?}", &data[0].value);
        } else {
            trace!("NODATA");
        }
        let data = serde_json::to_value(data)?;
        let req = ZabbixRequest::new(Self::HISTORY_DATA, &self.name, data);
        if let Ok(r) = self.send_request(&req, false) {
            if let ProxyResponse::RESPONSE(c) = r {
                trace!("{:?}", c);
                return Ok(c.success() && c.ok());
            }
        }
        Ok(false)
    }

    ///
    /// 返回服务器配置的代理主机和监控项
    ///
    /// compress = &[] 不压缩 item
    /// compress = &["["] 表示将 item 合并，去除 []
    /// compress = &["[", "_"] 表示将 item 合并，去除 [] 和 _
    ///
    pub fn get_proxy_config(&self, compress: &[&str]) -> Option<(HashSet<Host>, HashSet<Item>)> {
        if let Some(v) = self.get_config() {
            return Some((
                Host::from(get_item(&v["hosts"]["fields"], &v["hosts"]["data"])),
                Item::from(
                    get_item(&v["items"]["fields"], &v["items"]["data"]),
                    compress,
                ),
            ));
        }
        None
    }
}

fn get_item(field: &Value, data: &Value) -> Vec<HashMap<String, Value>> {
    let mut result = Vec::new();

    if let Some(field) = field.as_array() {
        let field = field.iter().map(|x| x.as_str());
        let data = data.as_array().unwrap();
        for x in data.iter() {
            let mut hm: HashMap<String, Value> = HashMap::new();
            let y = field.clone().zip(x.as_array().unwrap().iter());
            for z in y {
                hm.insert(z.0.unwrap().to_string(), z.1.clone());
            }
            result.push(hm);
        }
    }
    result
}