manycast/net/
mod.rs

1extern crate byteorder;
2use std::io::{Cursor, Read, Write};
3
4use crate::custom_module::manycastr::Address;
5use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
6use prost::bytes::Buf;
7
8pub(crate) mod packet;
9
10/// A struct detailing an IPv4Packet <https://en.wikipedia.org/wiki/Internet_Protocol_version_4>
11#[derive(Debug)]
12pub struct IPv4Packet {
13    pub length: u16,            // 16-bit Total Length
14    pub ttl: u8,                // 8-bit Time To Live
15    pub src: u32,               // 32-bit Source IP Address
16    pub dst: u32,               // 32-bit Destination IP Address
17    pub payload: PacketPayload, // Payload
18}
19
20/// Convert list of u8 (i.e. received bytes) into an IPv4Packet
21impl From<&[u8]> for IPv4Packet {
22    fn from(data: &[u8]) -> Self {
23        let mut cursor = Cursor::new(data);
24        // Get header length, which is the 4 right bits in the first byte (hence & 0xF)
25        // header length is in number of 32 bits i.e. 4 bytes (hence *4)
26        let header_length: usize = ((cursor.read_u8().unwrap() & 0xF) * 4).into(); // Total Length
27        cursor.set_position(8); // Time To Live
28        let ttl = cursor.read_u8().unwrap();
29        let packet_type = cursor.read_u8().unwrap(); // Protocol
30        cursor.set_position(12); // Address fields
31        let src = cursor.read_u32::<NetworkEndian>().unwrap(); // Source IP Address
32        let dst = cursor.read_u32::<NetworkEndian>().unwrap(); // Destination IP Address
33
34        // If the header length is longer than the data, the packet is incomplete
35        if header_length > data.len() {
36            return IPv4Packet {
37                length: header_length as u16,
38                ttl,
39                src,
40                dst,
41                payload: PacketPayload::Unimplemented,
42            };
43        }
44
45        let payload_bytes = &cursor.into_inner()[header_length..];
46        let payload = match packet_type {
47            1 => {
48                if payload_bytes.len() < 8 {
49                    PacketPayload::Unimplemented
50                } else {
51                    PacketPayload::Icmp {
52                        value: ICMPPacket::from(payload_bytes),
53                    }
54                }
55            }
56            17 => {
57                if payload_bytes.len() < 8 {
58                    PacketPayload::Unimplemented
59                } else {
60                    PacketPayload::Udp {
61                        value: UDPPacket::from(payload_bytes),
62                    }
63                }
64            }
65            6 => {
66                if payload_bytes.len() < 20 {
67                    PacketPayload::Unimplemented
68                } else {
69                    PacketPayload::Tcp {
70                        value: TCPPacket::from(payload_bytes),
71                    }
72                }
73            }
74            _ => PacketPayload::Unimplemented,
75        };
76
77        IPv4Packet {
78            length: header_length as u16,
79            ttl,
80            src,
81            dst,
82            payload,
83        }
84    }
85}
86
87/// Convert IPv4Packet into a vector of bytes
88impl From<&IPv4Packet> for Vec<u8> {
89    fn from(packet: &IPv4Packet) -> Self {
90        let (payload_type, payload) = match &packet.payload {
91            PacketPayload::Icmp { value } => (1, value.into()),
92            PacketPayload::Udp { value } => (17, value.into()),
93            PacketPayload::Tcp { value } => (6, value.into()),
94            PacketPayload::Unimplemented => (0, vec![]),
95        };
96
97        let mut wtr = vec![];
98        wtr.write_u8(0x45)
99            .expect("Unable to write to byte buffer for IPv4 packet"); // Version (4) and header length (5)
100        wtr.write_u8(0x00)
101            .expect("Unable to write to byte buffer for IPv4 packet"); // Type of Service
102        wtr.write_u16::<NetworkEndian>(packet.length)
103            .expect("Unable to write to byte buffer for IPv4 packet"); // Total Length
104        wtr.write_u16::<NetworkEndian>(0x3a7d)
105            .expect("Unable to write to byte buffer for IPv4 packet"); // Identification
106        wtr.write_u16::<NetworkEndian>(0x0000)
107            .expect("Unable to write to byte buffer for IPv4 packet"); // Flags (0) and Fragment Offset (0)
108        wtr.write_u8(packet.ttl)
109            .expect("Unable to write to byte buffer for IPv4 packet"); // Time To Live
110        wtr.write_u8(payload_type)
111            .expect("Unable to write to byte buffer for IPv4 packet"); // Protocol (ICMP)
112        wtr.write_u16::<NetworkEndian>(0x0000)
113            .expect("Unable to write to byte buffer for IPv4 packet"); // Header Checksum
114        wtr.write_u32::<NetworkEndian>(packet.src)
115            .expect("Unable to write to byte buffer for IPv4 packet"); // Source IP Address
116        wtr.write_u32::<NetworkEndian>(packet.dst)
117            .expect("Unable to write to byte buffer for IPv4 packet"); // Destination IP Address
118
119        // Calculate and write the checksum
120        let checksum = ICMPPacket::calc_checksum(&wtr); // Calculate checksum
121        let mut cursor = Cursor::new(wtr);
122        cursor.set_position(10); // Skip version (1 byte) and header length (1 byte)
123        cursor.write_u16::<NetworkEndian>(checksum).unwrap();
124
125        // Add the payload
126        cursor.set_position(20); // Skip the header
127        cursor
128            .write_all(&payload)
129            .expect("Unable to write to byte buffer for IPv4 packet"); // Payload
130
131        cursor.into_inner()
132    }
133}
134
135/// A struct detailing an IPv6Packet <https://en.wikipedia.org/wiki/IPv6>
136#[derive(Debug)]
137pub struct IPv6Packet {
138    // pub version: u8,                 // 4-bit Version
139    // pub traffic_class: u8,           // 8-bit Traffic Class
140    // pub flow_label: u32,             // 20-bit Flow Label
141    pub payload_length: u16,    // 16-bit Payload Length
142    pub next_header: u8,        // 8-bit Next Header
143    pub hop_limit: u8,          // 8-bit Hop Limit
144    pub src: u128,              // 128-bit Source Address
145    pub dst: u128,              // 128-bit Destination Address
146    pub payload: PacketPayload, // Payload
147}
148
149/// Convert bytes into an IPv6Packet
150impl From<&[u8]> for IPv6Packet {
151    fn from(data: &[u8]) -> Self {
152        let mut cursor = Cursor::new(data);
153        let _version_traffic_flow: u32 = cursor.read_u32::<NetworkEndian>().unwrap();
154        let payload_length = cursor.read_u16::<NetworkEndian>().unwrap();
155        let next_header = cursor.read_u8().unwrap();
156        let hop_limit = cursor.read_u8().unwrap();
157
158        let src = cursor.read_u128::<NetworkEndian>().unwrap(); // Source Address
159        let dst = cursor.read_u128::<NetworkEndian>().unwrap(); // Destination Address
160        let payload = &cursor.into_inner()[40..]; // IPv6 header is 40 bytes
161
162        // Implement PacketPayload based on the next_header value
163        let payload = match next_header {
164            58 => {
165                // ICMPv6
166                PacketPayload::Icmp {
167                    value: ICMPPacket::from(payload),
168                }
169            }
170            17 => {
171                // UDP
172                if payload.len() < 8 {
173                    PacketPayload::Unimplemented
174                } else {
175                    PacketPayload::Udp {
176                        value: UDPPacket::from(payload),
177                    }
178                }
179            }
180            6 => {
181                // TCP
182                if payload.len() < 20 {
183                    PacketPayload::Unimplemented
184                } else {
185                    PacketPayload::Tcp {
186                        value: TCPPacket::from(payload),
187                    }
188                }
189            }
190            _ => PacketPayload::Unimplemented, // Extension headers
191        };
192
193        IPv6Packet {
194            // version: (version_traffic_class >> 12) as u8,
195            // traffic_class: ((version_traffic_class >> 4) & 0xFF) as u8,
196            // flow_label,
197            payload_length,
198            next_header,
199            hop_limit,
200            src,
201            dst,
202            payload,
203        }
204    }
205}
206
207/// Convert an IPv6Packet into bytes
208impl From<&IPv6Packet> for Vec<u8> {
209    fn from(packet: &IPv6Packet) -> Self {
210        let mut wtr = vec![];
211        // Write traffic class 0x60 and flow label 0x003a7d
212        wtr.write_u32::<NetworkEndian>(0x60003a7d)
213            .expect("Unable to write to byte buffer for IPv6Packet");
214        wtr.write_u16::<NetworkEndian>(packet.payload_length)
215            .expect("Unable to write to byte buffer for IPv6Packet");
216        wtr.write_u8(packet.next_header)
217            .expect("Unable to write to byte buffer for IPv6Packet");
218        wtr.write_u8(packet.hop_limit)
219            .expect("Unable to write to byte buffer for IPv6Packet");
220        wtr.write_u128::<NetworkEndian>(packet.src)
221            .expect("Unable to write source address to byte buffer for IPv6Packet");
222        wtr.write_u128::<NetworkEndian>(packet.dst)
223            .expect("Unable to write destination address to byte buffer for IPv6Packet");
224
225        let payload = match &packet.payload {
226            PacketPayload::Icmp { value } => value.into(),
227            PacketPayload::Udp { value } => value.into(),
228            PacketPayload::Tcp { value } => value.into(),
229            PacketPayload::Unimplemented => vec![],
230        };
231
232        wtr.write_all(&payload)
233            .expect("Unable to write payload to byte buffer for IPv6Packet");
234
235        wtr
236    }
237}
238
239/// Definition of the IPV4Packet payload (either ICMPv4, UDP, TCP, or unimplemented)
240#[derive(Debug)]
241pub enum PacketPayload {
242    Icmp { value: ICMPPacket },
243    Udp { value: UDPPacket },
244    Tcp { value: TCPPacket },
245    Unimplemented,
246}
247
248/// Convert a packet payload to bytes
249impl From<PacketPayload> for Vec<u8> {
250    fn from(payload: PacketPayload) -> Self {
251        match payload {
252            PacketPayload::Icmp { value } => (&value).into(),
253            PacketPayload::Udp { value } => (&value).into(),
254            PacketPayload::Tcp { value } => (&value).into(),
255            PacketPayload::Unimplemented => vec![],
256        }
257    }
258}
259
260/// Struct defining the IPv4 pseudo-header for checksum calculation.
261#[derive(Debug)]
262pub struct PseudoHeaderV4 {
263    pub src: u32,
264    pub dst: u32,
265    pub protocol: u8,
266    pub length: u16,
267}
268
269/// Converting PsuedoHeader to bytes
270impl From<&PseudoHeaderV4> for Vec<u8> {
271    fn from(header: &PseudoHeaderV4) -> Self {
272        let mut wtr = vec![];
273        wtr.write_u32::<NetworkEndian>(header.src)
274            .expect("Unable to write to byte buffer for PseudoHeader");
275        wtr.write_u32::<NetworkEndian>(header.dst)
276            .expect("Unable to write to byte buffer for PseudoHeader");
277        wtr.write_u8(0) // 8 bits of zeroes
278            .expect("Unable to write to byte buffer for PseudoHeader");
279        wtr.write_u8(header.protocol)
280            .expect("Unable to write to byte buffer for PseudoHeader");
281        wtr.write_u16::<NetworkEndian>(header.length)
282            .expect("Unable to write to byte buffer for PseudoHeader");
283
284        wtr
285    }
286}
287
288/// Struct defining the IPv6 pseudo-header for checksum calculation.
289#[derive(Debug)]
290pub struct PseudoHeaderV6 {
291    pub src: u128,
292    pub dst: u128,
293    pub upper_layer_packet_length: u32,
294    pub next_header: u8,
295}
296
297/// Converting PsuedoHeaderv6 to bytes
298impl From<&PseudoHeaderV6> for Vec<u8> {
299    fn from(header: &PseudoHeaderV6) -> Self {
300        let mut wtr = vec![];
301        wtr.write_u128::<NetworkEndian>(header.src)
302            .expect("Unable to write to byte buffer for PseudoHeader");
303        wtr.write_u128::<NetworkEndian>(header.dst)
304            .expect("Unable to write to byte buffer for PseudoHeader");
305        wtr.write_u32::<NetworkEndian>(header.upper_layer_packet_length)
306            .expect("Unable to write to byte buffer for PseudoHeader");
307        wtr.write_u24::<NetworkEndian>(0) // zeroes
308            .expect("Unable to write to byte buffer for PseudoHeader");
309        wtr.write_u8(header.next_header)
310            .expect("Unable to write to byte buffer for PseudoHeader");
311
312        wtr
313    }
314}
315
316/// Struct defining a pseudo header that is used by both TCP and UDP to calculate their checksum
317#[derive(Debug)]
318pub enum PseudoHeader {
319    V4(PseudoHeaderV4),
320    V6(PseudoHeaderV6),
321}
322
323impl PseudoHeader {
324    /// Creates a new PseudoHeader (V4 or V6) based on the address types.
325    pub fn new(src_addr: &Address, dst_addr: &Address, protocol: u8, packet_length: u32) -> Self {
326        if src_addr.is_v6() {
327            Self::V6(PseudoHeaderV6 {
328                src: src_addr.get_v6(),
329                dst: dst_addr.get_v6(),
330                upper_layer_packet_length: packet_length, // Length field is 32 bits for IPv6
331                next_header: protocol,
332            })
333        } else {
334            Self::V4(PseudoHeaderV4 {
335                src: src_addr.get_v4(),
336                dst: dst_addr.get_v4(),
337                protocol,
338                length: packet_length as u16, // Length field is 16 bits for IPv4
339            })
340        }
341    }
342}
343
344/// Convert PseudoHeader to bytes
345impl From<&PseudoHeader> for Vec<u8> {
346    fn from(header: &PseudoHeader) -> Self {
347        match header {
348            PseudoHeader::V4(header) => header.into(),
349            PseudoHeader::V6(header) => header.into(),
350        }
351    }
352}
353
354/// Calculate the checksum for a UDP/TCP packet.
355///
356/// # Arguments
357///
358/// * 'buffer' - the UDP/TCP packet as bytes (without the IP header)
359///
360/// * 'pseudo_header' - the pseudo header for this packet (IPv4 or IPv6)
361pub fn calculate_checksum(buffer: &[u8], pseudo_header: &PseudoHeader) -> u16 {
362    let mut sum = 0u32;
363    let mut packet: Vec<u8> = pseudo_header.into();
364    packet.extend_from_slice(buffer);
365
366    // Sum the packet buffer
367    let packet_len = packet.len();
368    for chunk in packet.chunks_exact(2) {
369        let word = u16::from_be_bytes([chunk[0], chunk[1]]);
370        sum += u32::from(word);
371    }
372
373    // If the packet length is odd, add the last byte as a half-word (padded with 0)
374    if !packet_len.is_multiple_of(2) {
375        sum += u32::from(packet[packet_len - 1]) << 8;
376    }
377
378    // Fold the sum to 16 bits by adding the carry
379    while (sum >> 16) != 0 {
380        sum = (sum & 0xffff) + (sum >> 16);
381    }
382
383    // Return the one's complement of the sum
384    !(sum as u16)
385}
386
387/// An ICMP4Packet (ping packet) <https://en.wikipedia.org/wiki/Internet_Control_Message_Protocol#header_rest>
388#[derive(Debug)]
389pub struct ICMPPacket {
390    pub icmp_type: u8,
391    pub code: u8,
392    pub checksum: u16,
393    pub identifier: u16,
394    pub sequence_number: u16,
395    pub body: Vec<u8>,
396}
397
398/// Parsing from bytes to ICMP4Packet
399impl From<&[u8]> for ICMPPacket {
400    fn from(data: &[u8]) -> Self {
401        let mut data = Cursor::new(data);
402        ICMPPacket {
403            icmp_type: data.read_u8().unwrap(),
404            code: data.read_u8().unwrap(),
405            checksum: data.read_u16::<NetworkEndian>().unwrap(),
406            identifier: data.read_u16::<NetworkEndian>().unwrap(),
407            sequence_number: data.read_u16::<NetworkEndian>().unwrap(),
408            body: data.into_inner()[8..].to_vec(),
409        }
410    }
411}
412
413/// Convert ICMp4Packet into a vector of bytes
414impl From<&ICMPPacket> for Vec<u8> {
415    fn from(packet: &ICMPPacket) -> Self {
416        let mut wtr = vec![];
417        wtr.write_u8(packet.icmp_type)
418            .expect("Unable to write to byte buffer for ICMP packet");
419        wtr.write_u8(packet.code)
420            .expect("Unable to write to byte buffer for ICMP packet");
421        wtr.write_u16::<NetworkEndian>(packet.checksum)
422            .expect("Unable to write to byte buffer for ICMP packet");
423        wtr.write_u16::<NetworkEndian>(packet.identifier)
424            .expect("Unable to write to byte buffer for ICMP packet");
425        wtr.write_u16::<NetworkEndian>(packet.sequence_number)
426            .expect("Unable to write to byte buffer for ICMP packet");
427        wtr.write_all(&packet.body)
428            .expect("Unable to write to byte buffer for ICMP packet");
429        wtr
430    }
431}
432
433impl ICMPPacket {
434    /// Create a basic ICMPv4 ECHO_REQUEST (8.0) packet with checksum.
435    ///
436    /// # Arguments
437    ///
438    /// * 'identifier' - the identifier for the ICMP header
439    ///
440    /// * 'sequence_number' - the sequence number for the ICMP header
441    ///
442    /// * 'body' - the ICMP payload
443    ///
444    /// * 'src' - the source address of the packet
445    ///
446    /// * 'dst' - the destination address of the packet
447    ///
448    /// * 'ttl' - the time to live of the packet
449    ///
450    /// * 'info_url' - the URL to be added to the packet payload (e.g., opt-out URL)
451    pub fn echo_request(
452        identifier: u16,
453        sequence_number: u16,
454        body: Vec<u8>,
455        src: u32,
456        dst: u32,
457        ttl: u8,
458        info_url: &str,
459    ) -> Vec<u8> {
460        let body_len = body.len() as u16;
461        let mut packet = ICMPPacket {
462            icmp_type: 8,
463            code: 0,
464            checksum: 0,
465            identifier,
466            sequence_number,
467            body,
468        };
469
470        // Turn everything into a vec of bytes and calculate checksum
471        let mut icmp_bytes: Vec<u8> = (&packet).into();
472        icmp_bytes.extend(info_url.bytes());
473        packet.checksum = ICMPPacket::calc_checksum(&icmp_bytes);
474
475        let v4_packet = IPv4Packet {
476            length: 20 + 8 + body_len + info_url.len() as u16,
477            ttl,
478            src,
479            dst,
480            payload: PacketPayload::Icmp { value: packet },
481        };
482
483        let mut bytes: Vec<u8> = (&v4_packet).into();
484        bytes.extend(info_url.bytes());
485
486        bytes
487    }
488
489    /// Create an ICMPv6 echo request packet with checksum
490    ///
491    /// # Arguments
492    ///
493    /// * 'identifier' - the identifier for this packet
494    ///
495    /// * 'sequence_number' - the sequence number for this packet
496    ///
497    /// * 'body' - the payload of the packet
498    ///
499    /// * 'src' - the source address of the packet
500    ///
501    /// * 'dst' - the destination address of the packet
502    ///
503    /// * 'hop_limit' - the hop limit (TTL) of the packet
504    ///
505    /// * 'info_url' - URL encoded in packet payload (e.g., opt-out URL)
506    pub fn echo_request_v6(
507        identifier: u16,
508        sequence_number: u16,
509        body: Vec<u8>,
510        src: u128,
511        dst: u128,
512        hop_limit: u8,
513        info_url: &str,
514    ) -> Vec<u8> {
515        let body_len = body.len() as u16;
516        let mut packet = ICMPPacket {
517            icmp_type: 128,
518            code: 0,
519            checksum: 0,
520            identifier,
521            sequence_number,
522            body,
523        };
524        let icmp_bytes: Vec<u8> = (&packet).into();
525
526        // Append a pseudo header to the ICMP packet bytes
527        let mut psuedo_header: Vec<u8> = Vec::new();
528        psuedo_header
529            .write_u128::<NetworkEndian>(src)
530            .expect("Unable to write to byte buffer for PseudoHeader");
531        psuedo_header
532            .write_u128::<NetworkEndian>(dst)
533            .expect("Unable to write to byte buffer for PseudoHeader");
534        psuedo_header
535            .write_u32::<NetworkEndian>((8 + body_len + info_url.len() as u16) as u32) // ICMP length
536            .expect("Unable to write to byte buffer for PseudoHeader"); // Length of ICMP header + body
537        psuedo_header.write_u8(0).unwrap(); // zeroes
538        psuedo_header.write_u8(0).unwrap(); // zeroes
539        psuedo_header.write_u8(0).unwrap(); // zeroes
540        psuedo_header.write_u8(58).unwrap(); // next header (58 => ICMPv6)
541        psuedo_header.extend(icmp_bytes); // Add the ICMP packet bytes
542        psuedo_header.extend(info_url.bytes()); // Add the INFO_URL bytes
543        packet.checksum = ICMPPacket::calc_checksum(psuedo_header.as_slice()); // Calculate the checksum
544
545        let v6_packet = IPv6Packet {
546            payload_length: 8 + body_len + info_url.len() as u16, // ICMP header (8 bytes) + body length
547            next_header: 58,                                      // ICMPv6
548            hop_limit,
549            src,
550            dst,
551            payload: PacketPayload::Icmp { value: packet },
552        };
553
554        let mut bytes: Vec<u8> = (&v6_packet).into();
555        bytes.extend(info_url.bytes());
556
557        bytes
558    }
559
560    /// Calculate the ICMP Checksum.
561    ///
562    /// This calculation covers the entire ICMP  message (16-bit one's complement).
563    /// Works for both ICMPv4 and ICMPv6
564    fn calc_checksum(buffer: &[u8]) -> u16 {
565        let mut cursor = Cursor::new(buffer);
566        let mut sum: u32 = 0;
567        while let Ok(word) = cursor.read_u16::<NetworkEndian>() {
568            // Sum all 16-bit words
569            sum += u32::from(word);
570        }
571        if let Ok(byte) = cursor.read_u8() {
572            // If there is a byte left, sum it
573            sum += u32::from(byte);
574        }
575        while sum >> 16 > 0 {
576            sum = (sum & 0xffff) + (sum >> 16);
577        }
578        !sum as u16
579    }
580}
581
582/// An UDPPacket (UDP packet) <https://en.wikipedia.org/wiki/User_Datagram_Protocol>
583#[derive(Debug)]
584pub struct UDPPacket {
585    pub sport: u16,
586    pub dport: u16,
587    pub length: u16,
588    pub checksum: u16,
589    pub body: Vec<u8>,
590}
591
592/// Parsing from bytes into UDPPacket
593impl From<&[u8]> for UDPPacket {
594    fn from(data: &[u8]) -> Self {
595        let mut data = Cursor::new(data);
596        UDPPacket {
597            sport: data.read_u16::<NetworkEndian>().unwrap(),
598            dport: data.read_u16::<NetworkEndian>().unwrap(),
599            length: data.read_u16::<NetworkEndian>().unwrap(),
600            checksum: data.read_u16::<NetworkEndian>().unwrap(),
601            body: data.into_inner()[8..].to_vec(),
602        }
603    }
604}
605
606/// Convert UDPPacket into a vector of bytes
607impl From<&UDPPacket> for Vec<u8> {
608    fn from(packet: &UDPPacket) -> Self {
609        let mut wtr = vec![];
610        wtr.write_u16::<NetworkEndian>(packet.sport)
611            .expect("Unable to write to byte buffer for UDP packet");
612        wtr.write_u16::<NetworkEndian>(packet.dport)
613            .expect("Unable to write to byte buffer for UDP packet");
614        wtr.write_u16::<NetworkEndian>(packet.length)
615            .expect("Unable to write to byte buffer for UDP packet");
616        wtr.write_u16::<NetworkEndian>(packet.checksum)
617            .expect("Unable to write to byte buffer for UDP packet");
618        wtr.write_all(&packet.body)
619            .expect("Unable to write to byte buffer for UDP packet");
620
621        wtr
622    }
623}
624
625/// DNS request body
626#[allow(dead_code)]
627#[derive(Debug)]
628pub struct DNSRecord {
629    pub transaction_id: u16,
630    pub flags: u16,
631    pub questions: u16,
632    pub answer: u16,
633    pub authority: u16,
634    pub additional: u16,
635    pub domain: String,
636    pub record_type: u16,
637    pub class: u16,
638    pub body: Vec<u8>, // Possible answer sections
639}
640
641/// DNS answer body
642#[allow(dead_code)]
643#[derive(Debug)]
644pub struct DNSAnswer {
645    pub domain: String,
646    pub record_type: u16,
647    pub class: u16,
648    pub ttl: u32,
649    pub data_length: u16,
650    pub data: Vec<u8>,
651}
652
653/// DNS TXT data record
654#[allow(dead_code)]
655#[derive(Debug)]
656pub struct TXTRecord {
657    pub txt_length: u8,
658    pub txt: String,
659}
660
661/// Read a DNS name that is contained in a DNS response.
662/// Returns the domain name string of the A record reply.
663fn read_dns_name(data: &mut Cursor<&[u8]>) -> String {
664    let mut result = String::new();
665    loop {
666        if !data.has_remaining() {
667            break;
668        }
669        let label_len = data.read_u8().unwrap();
670        // If label length is 0, it is the end of the string
671        if label_len == 0 {
672            break;
673        }
674        // If the first two bytes of the label length is set to 11, it points to a different position
675        if label_len & 0xC0 == 0xC0 {
676            // The offset is the pointer to the previous domain name
677            let offset = ((label_len as u16 & 0x3F) << 8) | data.read_u8().unwrap() as u16;
678            data.set_position(offset as u64);
679            result.push_str(&read_dns_name(data));
680            break;
681        }
682        // Read the label
683        let mut label_bytes = vec![0; label_len as usize];
684
685        match data.read_exact(&mut label_bytes) {
686            Ok(()) => {}
687            Err(_) => {
688                return "Invalid domain name".to_string();
689            }
690        }
691
692        let label = String::from_utf8_lossy(&label_bytes).to_string();
693        result.push_str(&label);
694        result.push('.');
695    }
696    // Remove the trailing '.' if there is one
697    if result.ends_with('.') {
698        result.pop();
699    }
700    result
701}
702
703/// Parsing from bytes into a DNS A record
704impl From<&[u8]> for DNSRecord {
705    fn from(data: &[u8]) -> Self {
706        let mut data = Cursor::new(data);
707
708        let transaction_id = data.read_u16::<NetworkEndian>().unwrap();
709        let flags = data.read_u16::<NetworkEndian>().unwrap();
710        let questions = data.read_u16::<NetworkEndian>().unwrap();
711        let answer = data.read_u16::<NetworkEndian>().unwrap();
712        let authority = data.read_u16::<NetworkEndian>().unwrap();
713        let additional = data.read_u16::<NetworkEndian>().unwrap();
714        let domain = read_dns_name(&mut data);
715
716        let (record_type, class, body) = if data.remaining() >= 4 {
717            let record_type = data.read_u16::<NetworkEndian>().unwrap();
718            let class = data.read_u16::<NetworkEndian>().unwrap();
719            let body = data.clone().into_inner()[data.position() as usize..].to_vec();
720            (record_type, class, body)
721        } else {
722            let record_type = 0;
723            let class = 0;
724            let body = vec![];
725            (record_type, class, body)
726        };
727
728        DNSRecord {
729            transaction_id,
730            flags,
731            questions,
732            answer,
733            authority,
734            additional,
735            domain,
736            record_type,
737            class,
738            body,
739        }
740    }
741}
742
743/// Parsing from bytes into a DNS A record
744impl From<&[u8]> for DNSAnswer {
745    fn from(data: &[u8]) -> Self {
746        let mut data = Cursor::new(data);
747
748        // Make sure data has the required length
749        if data.remaining() < 10 {
750            return DNSAnswer {
751                domain: "Invalid DNS record".to_string(),
752                record_type: 0,
753                class: 0,
754                ttl: 0,
755                data_length: 0,
756                data: vec![],
757            };
758        }
759
760        DNSAnswer {
761            domain: data.read_u16::<NetworkEndian>().unwrap().to_string(), //read_dns_name(&mut data), // Two bytes that are a pointer to the domain name of the request record
762            record_type: data.read_u16::<NetworkEndian>().unwrap(),
763            class: data.read_u16::<NetworkEndian>().unwrap(),
764            ttl: data.read_u32::<NetworkEndian>().unwrap(),
765            data_length: data.read_u16::<NetworkEndian>().unwrap(),
766            data: data.clone().into_inner()[data.position() as usize..].to_vec(),
767        }
768    }
769}
770
771/// Parsing from bytes into a DNS TXT record
772impl From<&[u8]> for TXTRecord {
773    fn from(data: &[u8]) -> Self {
774        let mut data = Cursor::new(data);
775        // Make sure txt_length is not out of bounds
776        if data.remaining() < 1 {
777            return TXTRecord {
778                txt_length: 0,
779                txt: "Invalid TXT record".to_string(),
780            };
781        }
782
783        let txt_length = data.read_u8().unwrap();
784
785        // Make sure txt_length is not out of bounds
786        if txt_length as usize > data.remaining() {
787            return TXTRecord {
788                txt_length,
789                txt: "Invalid TXT record".to_string(),
790            };
791        }
792
793        TXTRecord {
794            txt_length,
795            // txt: read_dns_name(&mut data),
796            txt: String::from_utf8_lossy(&data.into_inner()[1..(1 + txt_length as u64) as usize])
797                .to_string(),
798        }
799    }
800}
801
802impl UDPPacket {
803    /// Create a UDP packet with a DNS A record request.
804    /// In the domain of the A record, we encode the transmit time, source and destination addresses, sender worker ID, and source port.
805    pub fn dns_request(
806        src: &Address,
807        dst: &Address,
808        sport: u16,
809        domain_name: &str,
810        tx_time: u64,
811        tx_id: u32,
812        ttl: u8,
813    ) -> Vec<u8> {
814        let dns_packet =
815            Self::create_a_record_request(domain_name, tx_time, src, dst, tx_id, sport);
816        let udp_length = (8 + dns_packet.len()) as u16;
817
818        let mut udp_packet = Self {
819            sport,
820            dport: 53u16, // DNS port
821            length: udp_length,
822            checksum: 0,
823            body: dns_packet,
824        };
825
826        let udp_bytes: Vec<u8> = (&udp_packet).into();
827
828        let pseudo_header = PseudoHeader::new(
829            src,
830            dst,
831            17, // UDP protocol
832            udp_length as u32,
833        );
834        udp_packet.checksum = calculate_checksum(&udp_bytes, &pseudo_header);
835
836        if src.is_v6() {
837            (&IPv6Packet {
838                payload_length: udp_length,
839                next_header: 17, // UDP
840                hop_limit: ttl,
841                src: src.get_v6(),
842                dst: dst.get_v6(),
843                payload: PacketPayload::Udp { value: udp_packet },
844            })
845                .into()
846        } else {
847            (&IPv4Packet {
848                length: 20 + udp_length,
849                ttl,
850                src: src.get_v4(),
851                dst: dst.get_v4(),
852                payload: PacketPayload::Udp { value: udp_packet },
853            })
854                .into()
855        }
856    }
857
858    /// Creating a DNS A Record Request body <http://www.tcpipguide.com/free/t_DNSMessageHeaderandQuestionSectionFormat.htm>
859    fn create_a_record_request(
860        domain_name: &str,
861        tx_time: u64,
862        src: &Address,
863        dst: &Address,
864        tx_id: u32,
865        sport: u16,
866    ) -> Vec<u8> {
867        // Max length of DNS domain name is 253 character
868        // Each label has a max length of 63 characters
869        // 20 + 10 + 10 + 3 + 5 + (4 '-' symbols) = 52 characters at most for subdomain
870        let subdomain = if src.is_v6() {
871            format!(
872                "{}.{}.{}.{}.{}.{}",
873                tx_time,
874                src.get_v6(),
875                dst.get_v6(),
876                tx_id,
877                sport,
878                domain_name
879            )
880        } else {
881            format!(
882                "{}.{}.{}.{}.{}.{}",
883                tx_time,
884                src.get_v4(),
885                dst.get_v4(),
886                tx_id,
887                sport,
888                domain_name
889            )
890        };
891        let mut dns_body: Vec<u8> = Vec::new();
892
893        // DNS Header
894        dns_body
895            .write_u16::<byteorder::BigEndian>(tx_id as u16)
896            .unwrap(); // Transaction ID
897        dns_body.write_u16::<byteorder::BigEndian>(0x0100).unwrap(); // Flags (Standard query, recursion desired)
898        dns_body.write_u16::<byteorder::BigEndian>(0x0001).unwrap(); // Number of questions
899        dns_body.write_u16::<byteorder::BigEndian>(0x0000).unwrap(); // Number of answer RRs
900        dns_body.write_u16::<byteorder::BigEndian>(0x0000).unwrap(); // Number of authority RRs
901        dns_body.write_u16::<byteorder::BigEndian>(0x0000).unwrap(); // Number of additional RRs
902
903        // DNS Question
904        for label in subdomain.split('.') {
905            dns_body.push(label.len() as u8);
906            dns_body.write_all(label.as_bytes()).unwrap();
907        }
908        dns_body.push(0); // Terminate the QNAME
909        dns_body.write_u16::<byteorder::BigEndian>(0x0001).unwrap(); // QTYPE (A record)
910        dns_body.write_u16::<byteorder::BigEndian>(0x0001).unwrap(); // QCLASS (IN)
911
912        dns_body
913    }
914
915    /// Create a UDP packet with a CHAOS TXT record request.
916    pub fn chaos_request(
917        src: &Address,
918        dst: &Address,
919        sport: u16,
920        tx: u32,
921        chaos: &str,
922    ) -> Vec<u8> {
923        let dns_body = Self::create_chaos_request(tx, chaos);
924        let udp_length = 8 + dns_body.len() as u32;
925
926        let mut udp_packet = Self {
927            sport,
928            dport: 53u16,
929            length: udp_length as u16,
930            checksum: 0,
931            body: dns_body,
932        };
933
934        let udp_bytes: Vec<u8> = (&udp_packet).into();
935
936        let pseudo_header = PseudoHeader::new(
937            src, dst, 17, // UDP protocol
938            udp_length,
939        );
940
941        udp_packet.checksum = calculate_checksum(&udp_bytes, &pseudo_header);
942
943        if src.is_v6() {
944            // Create the IPv6 packet
945            let v6_packet = IPv6Packet {
946                payload_length: udp_length as u16,
947                next_header: 17, // UDP
948                hop_limit: 255,
949                src: src.get_v6(),
950                dst: dst.get_v6(),
951                payload: PacketPayload::Udp { value: udp_packet },
952            };
953            (&v6_packet).into()
954        } else {
955            // Create the IPv4 packet
956            let v4_packet = IPv4Packet {
957                length: 20 + udp_length as u16,
958                ttl: 255,
959                src: src.get_v4(),
960                dst: dst.get_v4(),
961                payload: PacketPayload::Udp { value: udp_packet },
962            };
963            (&v4_packet).into()
964        }
965    }
966
967    /// Creating a DNS TXT record request body for id.orchestrator CHAOS request
968    fn create_chaos_request(tx_id: u32, chaos: &str) -> Vec<u8> {
969        let mut dns_body: Vec<u8> = Vec::new();
970
971        // DNS Header
972        dns_body.write_u32::<byteorder::BigEndian>(tx_id).unwrap(); // Transaction ID
973        dns_body.write_u16::<byteorder::BigEndian>(0x0100).unwrap(); // Flags (Standard query, recursion desired)
974        dns_body.write_u16::<byteorder::BigEndian>(0x0001).unwrap(); // Number of questions
975        dns_body.write_u16::<byteorder::BigEndian>(0x0000).unwrap(); // Number of answer RRs
976        dns_body.write_u16::<byteorder::BigEndian>(0x0000).unwrap(); // Number of authority RRs
977        dns_body.write_u16::<byteorder::BigEndian>(0x0000).unwrap(); // Number of additional RRs
978
979        // DNS Question (id.orchestrator)
980        for label in chaos.split('.') {
981            dns_body.push(label.len() as u8);
982            dns_body.write_all(label.as_bytes()).unwrap();
983        }
984        dns_body.push(0); // Terminate the QNAME
985        dns_body.write_u16::<byteorder::BigEndian>(0x0010).unwrap(); // QTYPE (TXT record)
986        dns_body.write_u16::<byteorder::BigEndian>(0x0003).unwrap(); // QCLASS (CHAOS)
987
988        dns_body
989    }
990}
991
992/// Get the length of a given domain in bytes
993/// TODO test this function
994#[allow(dead_code)]
995pub fn get_domain_bytes_length(domain: &str) -> u32 {
996    if domain == "." {
997        return 1;
998    }
999    let mut length = 1; // null terminator byte
1000    for label in domain.split('.') {
1001        length += label.len() as u32 + 1; // +1 for the label length byte
1002    }
1003    length
1004}
1005
1006/// A TCPPacket <https://en.wikipedia.org/wiki/Transmission_Control_Protocol>
1007#[derive(Debug)]
1008pub struct TCPPacket {
1009    pub sport: u16,
1010    pub dport: u16,
1011    pub seq: u32,
1012    pub ack: u32,
1013    // offset and reserved are combined into a single u8 (reserved is all 0's)
1014    pub offset: u8,
1015    pub flags: u8,
1016    pub window_size: u16,
1017    pub checksum: u16,
1018    pub pointer: u16,
1019    pub body: Vec<u8>,
1020}
1021
1022/// Parsing from bytes to TCPPacket
1023impl From<&[u8]> for TCPPacket {
1024    fn from(data: &[u8]) -> Self {
1025        let mut data = Cursor::new(data);
1026        TCPPacket {
1027            sport: data.read_u16::<NetworkEndian>().unwrap(),
1028            dport: data.read_u16::<NetworkEndian>().unwrap(),
1029            seq: data.read_u32::<NetworkEndian>().unwrap(),
1030            ack: data.read_u32::<NetworkEndian>().unwrap(),
1031            offset: data.read_u8().unwrap(),
1032            flags: data.read_u8().unwrap(),
1033            window_size: data.read_u16::<NetworkEndian>().unwrap(),
1034            checksum: data.read_u16::<NetworkEndian>().unwrap(),
1035            pointer: data.read_u16::<NetworkEndian>().unwrap(),
1036            body: data.into_inner()[8..].to_vec(),
1037        }
1038    }
1039}
1040
1041impl From<&TCPPacket> for Vec<u8> {
1042    fn from(packet: &TCPPacket) -> Self {
1043        let mut wtr = vec![];
1044
1045        wtr.write_u16::<NetworkEndian>(packet.sport)
1046            .expect("Unable to write to byte buffer for TCP packet");
1047        wtr.write_u16::<NetworkEndian>(packet.dport)
1048            .expect("Unable to write to byte buffer for TCP packet");
1049        wtr.write_u32::<NetworkEndian>(packet.seq)
1050            .expect("Unable to write to byte buffer for TCP packet");
1051        wtr.write_u32::<NetworkEndian>(packet.ack)
1052            .expect("Unable to write to byte buffer for TCP packet");
1053        wtr.write_u8(packet.offset)
1054            .expect("Unable to write to byte buffer for TCP packet");
1055        wtr.write_u8(packet.flags)
1056            .expect("Unable to write to byte buffer for TCP packet");
1057        wtr.write_u16::<NetworkEndian>(packet.window_size)
1058            .expect("Unable to write to byte buffer for TCP packet");
1059        wtr.write_u16::<NetworkEndian>(packet.checksum)
1060            .expect("Unable to write to byte buffer for TCP packet");
1061        wtr.write_u16::<NetworkEndian>(packet.pointer)
1062            .expect("Unable to write to byte buffer for TCP packet");
1063        wtr.write_all(&packet.body)
1064            .expect("Unable to write to byte buffer for TCP packet");
1065
1066        wtr
1067    }
1068}
1069
1070impl TCPPacket {
1071    /// Create a basic TCP SYN/ACK packet with checksum
1072    pub fn tcp_syn_ack(
1073        src: &Address,
1074        dst: &Address,
1075        sport: u16,
1076        dport: u16,
1077        ack: u32,
1078        ttl: u8,
1079        info_url: &str,
1080    ) -> Vec<u8> {
1081        let mut packet = Self {
1082            sport,
1083            dport,
1084            seq: 0, // Sequence number is not reflected
1085            ack,
1086            offset: 0b01010000, // Offset 5 for minimum TCP header length (0101) + 0000 for reserved
1087            flags: 0b00010010,  // SYN and ACK flags
1088            checksum: 0,
1089            pointer: 0,
1090            body: info_url.bytes().collect(),
1091            window_size: 0,
1092        };
1093
1094        // Turn everything into a vec of bytes and calculate checksum
1095        let bytes: Vec<u8> = (&packet).into();
1096
1097        let pseudo_header = PseudoHeader::new(
1098            src,
1099            dst,
1100            6,                  // TCP protocol
1101            bytes.len() as u32, // Length of the TCP header and data (measured in octets)
1102        );
1103        packet.checksum = calculate_checksum(&bytes, &pseudo_header);
1104
1105        if src.is_v6() {
1106            // Create the IPv6 packet
1107            let v6_packet = IPv6Packet {
1108                payload_length: bytes.len() as u16,
1109                next_header: 6, // TCP
1110                hop_limit: ttl,
1111                src: src.get_v6(),
1112                dst: dst.get_v6(),
1113                payload: PacketPayload::Tcp { value: packet },
1114            };
1115            (&v6_packet).into()
1116        } else {
1117            // Create the IPv4 packet
1118            let v4_packet = IPv4Packet {
1119                length: 20 + bytes.len() as u16,
1120                ttl,
1121                src: src.get_v4(),
1122                dst: dst.get_v4(),
1123                payload: PacketPayload::Tcp { value: packet },
1124            };
1125
1126            (&v4_packet).into()
1127        }
1128    }
1129}