mirror of
https://github.com/valmojr/armatak.git
synced 2026-06-13 20:24:34 +00:00
Added transport layer and configured extension commands to call mTLS socket connection
This commit is contained in:
153
src/tcp/mod.rs
153
src/tcp/mod.rs
@@ -1,118 +1,88 @@
|
||||
use arma_rs::Context;
|
||||
use lazy_static::lazy_static;
|
||||
use log::info;
|
||||
use std::io::Write;
|
||||
use std::net::TcpStream;
|
||||
use std::sync::mpsc::{self, Receiver, Sender};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::thread;
|
||||
|
||||
mod client;
|
||||
mod config;
|
||||
mod tls;
|
||||
mod transport;
|
||||
|
||||
pub mod cot;
|
||||
pub mod draw;
|
||||
|
||||
pub enum TcpCommand {
|
||||
SendMessage(String, Context),
|
||||
Stop,
|
||||
}
|
||||
|
||||
pub struct TcpClient {
|
||||
pub(crate) tx: Sender<TcpCommand>,
|
||||
}
|
||||
|
||||
impl TcpClient {
|
||||
pub fn start(&self, address: String, rx: Receiver<TcpCommand>, ctx: Context) {
|
||||
if let Some(ref client) = *TCP_CLIENT.lock().unwrap() {
|
||||
client.stop();
|
||||
}
|
||||
|
||||
let connection = Arc::new(Mutex::new(None));
|
||||
let connection_clone = Arc::clone(&connection);
|
||||
|
||||
thread::spawn(move || {
|
||||
let mut running = true;
|
||||
|
||||
let tcp_thread = thread::spawn(move || match TcpStream::connect(&address) {
|
||||
Ok(stream) => {
|
||||
let _ = ctx.callback_data("TCP SOCKET", "Connected to TCP Server", address);
|
||||
*connection_clone.lock().unwrap() = Some(stream);
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = ctx.callback_data(
|
||||
"TCP SOCKET ERROR",
|
||||
"TAK Socket connection failed",
|
||||
e.to_string(),
|
||||
);
|
||||
info!("Failed to connect to TCP server: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
while running {
|
||||
match rx.recv() {
|
||||
Ok(TcpCommand::SendMessage(message, context)) => {
|
||||
if let Some(mut stream) = connection.lock().unwrap().as_ref() {
|
||||
if let Err(e) = stream.write_all(message.as_bytes()) {
|
||||
info!("Failed to send message: {}", e);
|
||||
|
||||
let _ = context.callback_data(
|
||||
"TCP SOCKET ERROR",
|
||||
"TAK Socket disconnected",
|
||||
e.to_string(),
|
||||
);
|
||||
|
||||
running = false;
|
||||
}
|
||||
} else {
|
||||
let _ = context.callback_null(
|
||||
"TCP SOCKET ERROR",
|
||||
"TAK Socket is not active",
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(TcpCommand::Stop) => {
|
||||
running = false;
|
||||
info!("Stopping TCP client.");
|
||||
}
|
||||
Err(error) => {
|
||||
info!("Error receiving command: {}", error.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tcp_thread.join().unwrap();
|
||||
});
|
||||
}
|
||||
|
||||
pub fn send_payload(&self, context: Context, payload: String) {
|
||||
let tx = self.tx.clone();
|
||||
thread::spawn(move || {
|
||||
tx.send(TcpCommand::SendMessage(payload, context)).unwrap();
|
||||
});
|
||||
}
|
||||
|
||||
pub fn stop(&self) {
|
||||
let tx = self.tx.clone();
|
||||
thread::spawn(move || {
|
||||
tx.send(TcpCommand::Stop).unwrap();
|
||||
});
|
||||
}
|
||||
}
|
||||
use client::{TcpClient, TcpCommand};
|
||||
use config::ConnectionConfig;
|
||||
use tls::artifacts::clear_enrollment_artifacts;
|
||||
|
||||
lazy_static! {
|
||||
static ref TCP_CLIENT: Arc<Mutex<Option<TcpClient>>> = Arc::new(Mutex::new(None));
|
||||
}
|
||||
|
||||
pub fn start(ctx: Context, address: String) -> &'static str {
|
||||
fn start_with_config(ctx: Context, config: ConnectionConfig) {
|
||||
let (tx, rx): (Sender<TcpCommand>, Receiver<TcpCommand>) = mpsc::channel();
|
||||
|
||||
let client = TcpClient { tx };
|
||||
client.start(address, rx, ctx);
|
||||
client.start(config, rx, ctx);
|
||||
|
||||
let mut client_guard = TCP_CLIENT.lock().unwrap();
|
||||
*client_guard = Some(client);
|
||||
}
|
||||
|
||||
pub fn start(ctx: Context, address: String) -> &'static str {
|
||||
start_with_config(ctx, ConnectionConfig::Plain { address });
|
||||
|
||||
"Starting TCP Client"
|
||||
}
|
||||
|
||||
pub fn start_mtls(
|
||||
ctx: Context,
|
||||
address: String,
|
||||
server_name: String,
|
||||
ca_cert_path: String,
|
||||
client_cert_path: String,
|
||||
client_key_path: String,
|
||||
) -> &'static str {
|
||||
start_with_config(
|
||||
ctx,
|
||||
ConnectionConfig::Mtls {
|
||||
address,
|
||||
server_name,
|
||||
ca_cert_path,
|
||||
client_cert_path,
|
||||
client_key_path,
|
||||
},
|
||||
);
|
||||
|
||||
"Starting mTLS TCP Client"
|
||||
}
|
||||
|
||||
pub fn start_enroll_mtls(
|
||||
ctx: Context,
|
||||
host: String,
|
||||
server_name: String,
|
||||
enroll_port: String,
|
||||
username: String,
|
||||
password: String,
|
||||
client_uid: String,
|
||||
) -> &'static str {
|
||||
clear_enrollment_artifacts();
|
||||
start_with_config(
|
||||
ctx,
|
||||
ConnectionConfig::EnrollMtls {
|
||||
host,
|
||||
server_name,
|
||||
enroll_port,
|
||||
username,
|
||||
password,
|
||||
client_uid,
|
||||
},
|
||||
);
|
||||
|
||||
"Starting enrolled mTLS TCP Client"
|
||||
}
|
||||
|
||||
pub fn send_payload(ctx: Context, payload: String) -> &'static str {
|
||||
if let Some(ref client) = *TCP_CLIENT.lock().unwrap() {
|
||||
client.send_payload(ctx, payload);
|
||||
@@ -128,6 +98,7 @@ pub fn stop(ctx: Context) -> &'static str {
|
||||
if let Some(ref client) = *TCP_CLIENT.lock().unwrap() {
|
||||
client.stop();
|
||||
let _ = ctx.callback_null("TCP SOCKET", "TCP client stopped");
|
||||
clear_enrollment_artifacts();
|
||||
} else {
|
||||
let _ = ctx.callback_null("TCP SOCKET ERROR", "TCP client is not running");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user