shohei 2.5.0

Infrastructure diagnostics library: DNS, DNSSEC, TLS certificate inspection, email security, DNS propagation, and MCP-integrated AI agent support
Documentation 
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
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254

//! DNS amplification attack potential checker — measure response size ratio.

use serde::{Deserialize, Serialize};
use std::net::{IpAddr, SocketAddr};
use std::str::FromStr;
use tokio::net::UdpSocket;
use tokio::time::{timeout, Duration};
use crate::error::Result;

/// Check DNS amplification potential for a nameserver.
pub async fn check_dns_amplification(req: &DnsAmplificationRequest) -> Result<DnsAmplificationResult> {
    let nameserver = &req.nameserver;
    let port = req.port;
    let domain = &req.domain;
    let timeout_secs = Duration::from_secs(req.timeout_secs);

    // Build DNS query for domain (A record)
    let query = build_dns_query(domain, 1)?;  // 1 = A record type
    let query_size = query.len();

    // Send UDP query to nameserver
    let socket = match UdpSocket::bind("0.0.0.0:0").await {
        Ok(s) => s,
        Err(e) => {
            return Ok(DnsAmplificationResult {
                nameserver: nameserver.clone(),
                domain: domain.clone(),
                query_size: query_size as u32,
                response_size: 0,
                amplification_factor: 0.0,
                risk_level: "error".to_string(),
                error: Some(format!("Socket bind failed: {}", e)),
            });
        }
    };

    // Validate nameserver is a parseable IP address and not a private/reserved address
    let ip = match IpAddr::from_str(nameserver) {
        Ok(ip) => ip,
        Err(_) => {
            return Ok(DnsAmplificationResult {
                nameserver: nameserver.clone(),
                domain: domain.clone(),
                query_size: query_size as u32,
                response_size: 0,
                amplification_factor: 0.0,
                risk_level: "error".to_string(),
                error: Some("Invalid IP address".to_string()),
            });
        }
    };
    if crate::api::helpers::is_private_or_special_ip(&ip) {
        return Ok(DnsAmplificationResult {
            nameserver: nameserver.clone(),
            domain: domain.clone(),
            query_size: query_size as u32,
            response_size: 0,
            amplification_factor: 0.0,
            risk_level: "error".to_string(),
            error: Some("Private or reserved IP address blocked".to_string()),
        });
    }
    let addr = SocketAddr::new(ip, port);
    if let Err(e) = socket.connect(addr).await {
        return Ok(DnsAmplificationResult {
            nameserver: nameserver.clone(),
            domain: domain.clone(),
            query_size: query_size as u32,
            response_size: 0,
            amplification_factor: 0.0,
            risk_level: "error".to_string(),
            error: Some(format!("Connect failed: {}", e)),
        });
    }

    // Send query
    if let Err(e) = socket.send(&query).await {
        return Ok(DnsAmplificationResult {
            nameserver: nameserver.clone(),
            domain: domain.clone(),
            query_size: query_size as u32,
            response_size: 0,
            amplification_factor: 0.0,
            risk_level: "error".to_string(),
            error: Some(format!("Send failed: {}", e)),
        });
    }

    // Receive response
    let mut buf = vec![0; 4096];  // DNS responses up to 4KB (UDP limit)
    let response_size = match timeout(timeout_secs, socket.recv(&mut buf)).await {
        Ok(Ok(size)) => size,
        Ok(Err(e)) => {
            return Ok(DnsAmplificationResult {
                nameserver: nameserver.clone(),
                domain: domain.clone(),
                query_size: query_size as u32,
                response_size: 0,
                amplification_factor: 0.0,
                risk_level: "error".to_string(),
                error: Some(format!("Receive failed: {}", e)),
            });
        }
        Err(_) => {
            return Ok(DnsAmplificationResult {
                nameserver: nameserver.clone(),
                domain: domain.clone(),
                query_size: query_size as u32,
                response_size: 0,
                amplification_factor: 0.0,
                risk_level: "error".to_string(),
                error: Some("Timeout".to_string()),
            });
        }
    };

    // Calculate amplification factor
    let amplification_factor = response_size as f64 / query_size as f64;

    // Risk assessment
    let risk_level = if amplification_factor > 50.0 {
        "critical".to_string()
    } else if amplification_factor > 20.0 {
        "high".to_string()
    } else if amplification_factor > 10.0 {
        "medium".to_string()
    } else if amplification_factor > 1.0 {
        "low".to_string()
    } else {
        "none".to_string()
    };

    Ok(DnsAmplificationResult {
        nameserver: nameserver.clone(),
        domain: domain.clone(),
        query_size: query_size as u32,
        response_size: response_size as u32,
        amplification_factor,
        risk_level,
        error: None,
    })
}

fn build_dns_query(domain: &str, record_type: u8) -> Result<Vec<u8>> {
    // Validate domain name before constructing query
    if domain.is_empty() {
        return Err(crate::error::ShoheError::Parse("Domain cannot be empty".to_string()));
    }
    if domain.len() > 253 {
        return Err(crate::error::ShoheError::Parse("Domain length exceeds 253 bytes".to_string()));
    }

    // Validate domain labels
    for label in domain.split('.') {
        if label.is_empty() {
            return Err(crate::error::ShoheError::Parse("Domain contains empty label".to_string()));
        }
        if label.len() > 63 {
            return Err(crate::error::ShoheError::Parse("Label exceeds 63 bytes".to_string()));
        }
        // Validate label characters: only alphanumeric, hyphen, and underscore
        if !label.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_') {
            return Err(crate::error::ShoheError::Parse("Label contains invalid characters".to_string()));
        }
        // Labels cannot start or end with hyphen
        if label.starts_with('-') || label.ends_with('-') {
            return Err(crate::error::ShoheError::Parse("Label cannot start or end with hyphen".to_string()));
        }
    }

    // Simplified DNS query builder (valid minimal DNS query)
    let mut query = Vec::new();

    // Transaction ID (2 bytes): monotonic counter for uniqueness
    static QUERY_ID: std::sync::atomic::AtomicU16 = std::sync::atomic::AtomicU16::new(1);
    let qid = QUERY_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    query.push((qid >> 8) as u8);
    query.push((qid & 0xff) as u8);

    // Flags (2 bytes): standard query (0x0000)
    query.push(0x00);
    query.push(0x00);

    // Question count (2 bytes): 1
    query.push(0x00);
    query.push(0x01);

    // Answer count (2 bytes): 0
    query.push(0x00);
    query.push(0x00);

    // Authority count (2 bytes): 0
    query.push(0x00);
    query.push(0x00);

    // Additional count (2 bytes): 0
    query.push(0x00);
    query.push(0x00);

    // Question section: encode domain name
    for label in domain.split('.') {
        query.push(label.len() as u8);
        query.extend_from_slice(label.as_bytes());
    }
    query.push(0x00);  // Root label

    // Question type (2 bytes): record_type
    query.push(0x00);
    query.push(record_type);

    // Question class (2 bytes): IN (0x0001)
    query.push(0x00);
    query.push(0x01);

    Ok(query)
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DnsAmplificationRequest {
    /// Nameserver IP address (IPv4 or IPv6)
    pub nameserver: String,
    /// Port (default 53)
    #[serde(default = "default_port")]
    pub port: u16,
    /// Domain to query
    #[serde(default = "default_domain")]
    pub domain: String,
    #[serde(default = "default_timeout")]
    pub timeout_secs: u64,
}

fn default_port() -> u16 {
    53
}

fn default_domain() -> String {
    "example.com".to_string()
}

fn default_timeout() -> u64 {
    crate::api::helpers::DEFAULT_REQUEST_TIMEOUT_SECS
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DnsAmplificationResult {
    pub nameserver: String,
    pub domain: String,
    pub query_size: u32,
    pub response_size: u32,
    pub amplification_factor: f64,
    /// "none", "low", "medium", "high", "critical", "error"
    pub risk_level: String,
    pub error: Option<String>,
}