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

use crate::error::ProtocolError;
use bytes::Bytes;
use uuid::Uuid;

use crate::message::payload::tcp::{AgentTcpPayload, ProxyTcpPayload};
use crate::message::payload::{
    tcp::ProxyTcpInitResult,
    udp::{AgentUdpPayload, ProxyUdpPayload},
};
use crate::message::values::address::PpaassUnifiedAddress;
use crate::message::values::encryption::PpaassMessagePayloadEncryption;
use crate::message::{
    PpaassAgentMessage, PpaassAgentMessagePayload, PpaassProxyMessage, PpaassProxyMessagePayload,
};

pub struct PpaassMessageGenerator;

impl PpaassMessageGenerator {
    /// Generate the agent tcp init message
    pub fn generate_agent_tcp_init_message(
        user_token: String,
        src_address: PpaassUnifiedAddress,
        dst_address: PpaassUnifiedAddress,
        encryption: PpaassMessagePayloadEncryption,
    ) -> Result<PpaassAgentMessage, ProtocolError> {
        let payload = PpaassAgentMessagePayload::Tcp(AgentTcpPayload::Init {
            src_address,
            dst_address,
        });
        let message =
            PpaassAgentMessage::new(Uuid::new_v4().to_string(), user_token, encryption, payload);
        Ok(message)
    }

    /// Generate the proxy tcp init message
    pub fn generate_proxy_tcp_init_message(
        user_token: String,
        src_address: PpaassUnifiedAddress,
        dst_address: PpaassUnifiedAddress,
        encryption: PpaassMessagePayloadEncryption,
        result: ProxyTcpInitResult,
    ) -> Result<PpaassProxyMessage, ProtocolError> {
        let payload = PpaassProxyMessagePayload::Tcp(ProxyTcpPayload::Init {
            src_address,
            dst_address,
            result,
        });
        let message = PpaassProxyMessage::new(
            Uuid::new_v4().to_string(),
            user_token.to_string(),
            encryption,
            payload,
        );
        Ok(message)
    }

    /// Generate the agent tcp data message
    pub fn generate_agent_tcp_data_message(
        user_token: String,
        encryption: PpaassMessagePayloadEncryption,
        data: Bytes,
    ) -> Result<PpaassAgentMessage, ProtocolError> {
        let payload = PpaassAgentMessagePayload::Tcp(AgentTcpPayload::Data { content: data });
        let message = PpaassAgentMessage::new(
            Uuid::new_v4().to_string(),
            user_token.to_string(),
            encryption,
            payload,
        );
        Ok(message)
    }

    /// Generate the proxy tcp data message
    pub fn generate_proxy_tcp_data_message(
        user_token: String,
        encryption: PpaassMessagePayloadEncryption,
        data: Bytes,
    ) -> Result<PpaassProxyMessage, ProtocolError> {
        let payload = PpaassProxyMessagePayload::Tcp(ProxyTcpPayload::Data { content: data });
        let message = PpaassProxyMessage::new(
            Uuid::new_v4().to_string(),
            user_token.to_string(),
            encryption,
            payload,
        );
        Ok(message)
    }

    /// Generate the agent udp data message
    pub fn generate_agent_udp_data_message(
        user_token: String,
        encryption: PpaassMessagePayloadEncryption,
        src_address: PpaassUnifiedAddress,
        dst_address: PpaassUnifiedAddress,
        data: Bytes,
        need_response: bool,
    ) -> Result<PpaassAgentMessage, ProtocolError> {
        let payload = PpaassAgentMessagePayload::Udp(AgentUdpPayload {
            src_address,
            dst_address,
            data,
            need_response,
        });
        let message = PpaassAgentMessage::new(
            Uuid::new_v4().to_string(),
            user_token.to_string(),
            encryption,
            payload,
        );
        Ok(message)
    }

    /// Generate the proxy udp data message
    pub fn generate_proxy_udp_data_message(
        user_token: String,
        encryption: PpaassMessagePayloadEncryption,
        src_address: PpaassUnifiedAddress,
        dst_address: PpaassUnifiedAddress,
        data: Bytes,
    ) -> Result<PpaassProxyMessage, ProtocolError> {
        let payload = PpaassProxyMessagePayload::Udp(ProxyUdpPayload {
            src_address,
            dst_address,
            data,
        });
        let message = PpaassProxyMessage::new(
            Uuid::new_v4().to_string(),
            user_token.to_string(),
            encryption,
            payload,
        );
        Ok(message)
    }
}