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

use stringprep::saslprep;

pub use super::errors::{IntegrityKeyGenerationError, MessageDecodeError, MessageEncodeError};
use crate::definitions::StunTransactionId;
use crate::header::StunHeader;
use crate::header::{StunMessageClass, StunMessageMethod};
use crate::StunAttribute;

use super::message::StunMessage;

impl StunMessage {
    /// Creates a new message
    pub fn new(method: StunMessageMethod, class: StunMessageClass) -> Self {
        let header = StunHeader::new(method, class, None);

        Self {
            header,
            attributes: Vec::new(),
        }
    }

    /// Creates a Binding Request
    pub fn create_request() -> Self {
        Self::default().set_message_class(StunMessageClass::Request)
    }

    /// Creates a Binding Success Response
    pub fn create_success_response() -> Self {
        Self::default().set_message_class(StunMessageClass::SuccessResponse)
    }

    /// Creates a Binding Error Response
    pub fn create_error_response() -> Self {
        Self::default().set_message_class(StunMessageClass::ErrorResponse)
    }

    /// Creates a Binding Indication
    pub fn create_indication() -> Self {
        Self::default().set_message_class(StunMessageClass::Indication)
    }

    /// Sets message transaction id
    pub fn set_transaction_id(mut self, transaction_id: StunTransactionId) -> Self {
        self.header.transaction_id = transaction_id;

        self
    }

    /// Sets message class
    pub fn set_message_class(mut self, class: StunMessageClass) -> Self {
        self.header.message_class = class;

        self
    }

    /// Sets message method
    pub fn set_message_method(mut self, method: StunMessageMethod) -> Self {
        self.header.message_method = method;

        self
    }

    /// Returns an immutable reference to the message header
    pub fn get_header(&self) -> &StunHeader {
        &self.header
    }

    /// Returns an immutable reference to the message attributes
    pub fn get_attributes(&self) -> &Vec<StunAttribute> {
        &self.attributes
    }

    /// Adds an attribute to the list
    pub fn add_attribute(mut self, attr: StunAttribute) -> Self {
        self.attributes.push(attr);

        self
    }

    /// Adds a Fingerprint attribute at the end of the message
    ///
    /// NOTE: This function should be invoked only when all other attributes are added
    pub fn add_fingerprint(mut self) -> Self {
        self.attributes
            .push(StunAttribute::Fingerprint { value: 0 });

        self
    }

    /// Adds a MessageIntegrity attribute at the end of the message
    ///
    /// NOTE: This function should be invoked only when all other attributes are added but before the Fingerprint attribute
    pub fn add_message_integrity(mut self) -> Self {
        self.attributes
            .push(StunAttribute::MessageIntegrity { key: Vec::new() });

        self
    }

    /// Adds USER, REALM and MESSAGE-INTEGRITY attributes for long term credential authentication
    ///
    /// NOTE: This function should be invoked only when all other attributes are added but before the Fingerprint attribute
    pub fn add_long_term_credential_message_integrity(
        mut self,
        username: &str,
        realm: &str,
    ) -> Result<Self, stringprep::Error> {
        self.attributes.push(StunAttribute::Username {
            value: saslprep(username)?.to_string(),
        });

        self.attributes.push(StunAttribute::Realm {
            value: saslprep(realm)?.to_string(),
        });

        Ok(self.add_message_integrity())
    }
}

impl std::default::Default for StunMessage {
    /// Default STUN message.
    ///
    /// Class: Request
    /// Method: Binding
    /// Transaction ID: randomly generated
    fn default() -> Self {
        Self::new(StunMessageMethod::BindingRequest, StunMessageClass::Request)
    }
}