plazmoid
4 years ago
parent
3e0c9ecd77
commit
4e88f49f96
25 changed files with 410 additions and 275 deletions
@ -1,3 +1,3 @@ |
||||
FROM centos:7 |
||||
|
||||
CMD yum update |
||||
RUN yum update -y |
||||
@ -1 +1,5 @@ |
||||
pub mod client; |
||||
pub mod panel; |
||||
|
||||
pub use client::AgentClient; |
||||
pub use panel::Panel; |
||||
|
||||
@ -1,3 +1,40 @@ |
||||
const BINARY: &str = "/u_panel"; |
||||
use serde_json::{from_slice, Value}; |
||||
use shlex::split; |
||||
use std::process::{Command, Output}; |
||||
|
||||
const PANEL_BINARY: &str = "/u_panel"; |
||||
|
||||
pub struct Panel; |
||||
|
||||
impl Panel { |
||||
fn run(args: &[&str]) -> Output { |
||||
Command::new(PANEL_BINARY) |
||||
.arg("--json") |
||||
.args(args) |
||||
.output() |
||||
.unwrap() |
||||
} |
||||
|
||||
pub fn output_argv(args: &[&str]) -> Value { |
||||
let result = Self::run(args); |
||||
assert!(result.status.success()); |
||||
from_slice(&result.stdout).unwrap() |
||||
} |
||||
|
||||
pub fn output(args: &str) -> Value { |
||||
let splitted = split(args).unwrap(); |
||||
Self::output_argv( |
||||
splitted |
||||
.iter() |
||||
.map(|s| s.as_ref()) |
||||
.collect::<Vec<&str>>() |
||||
.as_ref(), |
||||
) |
||||
} |
||||
|
||||
pub fn check_output(args: &str) -> Vec<Value> { |
||||
let result = Self::output(args); |
||||
assert_eq!(result["status"], "ok"); |
||||
result["data"].as_array().unwrap().clone() |
||||
} |
||||
} |
||||
|
||||
@ -1,2 +1,54 @@ |
||||
mod behaviour; |
||||
mod helpers; |
||||
|
||||
use helpers::{AgentClient, Panel}; |
||||
|
||||
use serde_json::json; |
||||
use uuid::Uuid; |
||||
|
||||
type TestResult<R = ()> = Result<R, Box<dyn std::error::Error>>; |
||||
|
||||
async fn register_agent() -> Uuid { |
||||
let cli = AgentClient::new(); |
||||
let agent_uid = Uuid::new_v4(); |
||||
let resp = cli.get(format!("get_agent_jobs/{}", agent_uid)).await; |
||||
let job_id = &resp["job_id"]; |
||||
let resp = cli.get(format!("get_jobs/{}", job_id)).await; |
||||
assert_eq!(&resp["alias"], "agent_hello"); |
||||
let agent_data = json! { |
||||
{"id": &agent_uid,"inner":[ |
||||
{"Agent": |
||||
{"alias":null, |
||||
"hostname":"3b1030fa6324", |
||||
"id":&agent_uid, |
||||
"is_root":false, |
||||
"is_root_allowed":false, |
||||
"last_active":{"secs_since_epoch":1625271265,"nanos_since_epoch":92814921}, |
||||
"platform":"x86_64-unknown-linux-gnu", |
||||
"regtime":{"secs_since_epoch":1625271265,"nanos_since_epoch":92814945}, |
||||
"state":"New", |
||||
"token":null, |
||||
"username":"root"} |
||||
}]} |
||||
}; |
||||
cli.post("report", &agent_data).await; |
||||
agent_uid |
||||
} |
||||
|
||||
#[tokio::test] |
||||
async fn test_first_connection() -> TestResult { |
||||
let uid = register_agent().await; |
||||
let agents = Panel::check_output("agents list"); |
||||
dbg!(&agents); |
||||
assert_eq!(agents.len(), 2); |
||||
let found = agents |
||||
.iter() |
||||
.find(|v| v["id"].as_str().unwrap() == uid.to_string()); |
||||
assert!(found.is_some()); |
||||
Ok(()) |
||||
} |
||||
|
||||
#[tokio::test] |
||||
async fn test_setup_tasks() -> TestResult { |
||||
register_agent().await; |
||||
Ok(()) |
||||
} |
||||
|
||||
@ -1,168 +0,0 @@ |
||||
use crate::UError; |
||||
use chrono::{offset::Local, DateTime}; |
||||
use nix::{ |
||||
sys::signal::{signal, SigHandler, Signal}, |
||||
unistd::{chdir, close as fdclose, fork, getppid, setsid, ForkResult}, |
||||
}; |
||||
use std::{ |
||||
env::temp_dir, fs, ops::Drop, os::unix::fs::PermissionsExt, path::PathBuf, process::exit, |
||||
time::SystemTime, |
||||
}; |
||||
use uuid::Uuid; |
||||
|
||||
pub trait OneOrMany<T> { |
||||
fn into_vec(self) -> Vec<T>; |
||||
} |
||||
|
||||
impl<T> OneOrMany<T> for T { |
||||
fn into_vec(self) -> Vec<T> { |
||||
vec![self] |
||||
} |
||||
} |
||||
|
||||
impl<T> OneOrMany<T> for Vec<T> { |
||||
fn into_vec(self) -> Vec<T> { |
||||
self |
||||
} |
||||
} |
||||
|
||||
pub struct TempFile { |
||||
path: PathBuf, |
||||
} |
||||
|
||||
impl TempFile { |
||||
pub fn get_path(&self) -> String { |
||||
self.path.to_string_lossy().to_string() |
||||
} |
||||
|
||||
pub fn new() -> Self { |
||||
let name = Uuid::simple(&Uuid::new_v4()).to_string(); |
||||
let mut path = temp_dir(); |
||||
path.push(name); |
||||
Self { path } |
||||
} |
||||
|
||||
pub fn write_all(&self, data: &[u8]) -> Result<(), String> { |
||||
fs::write(&self.path, data).map_err(|e| e.to_string()) |
||||
} |
||||
|
||||
pub fn write_exec(data: &[u8]) -> Result<Self, (String, String)> { |
||||
let this = Self::new(); |
||||
let path = this.get_path(); |
||||
this.write_all(data).map_err(|e| (path.clone(), e))?; |
||||
let perms = fs::Permissions::from_mode(0o555); |
||||
fs::set_permissions(&path, perms).map_err(|e| (path, e.to_string()))?; |
||||
Ok(this) |
||||
} |
||||
} |
||||
|
||||
impl Drop for TempFile { |
||||
fn drop(&mut self) { |
||||
fs::remove_file(&self.path).ok(); |
||||
} |
||||
} |
||||
|
||||
pub struct CombinedResult<T, E = UError> { |
||||
ok: Vec<T>, |
||||
err: Vec<E>, |
||||
} |
||||
|
||||
impl<T, E> CombinedResult<T, E> { |
||||
pub fn new() -> Self { |
||||
Self { |
||||
ok: vec![], |
||||
err: vec![], |
||||
} |
||||
} |
||||
|
||||
pub fn ok<I: OneOrMany<T>>(&mut self, result: I) { |
||||
self.ok.extend(result.into_vec()); |
||||
} |
||||
|
||||
pub fn err<I: OneOrMany<E>>(&mut self, err: I) { |
||||
self.err.extend(err.into_vec()); |
||||
} |
||||
|
||||
pub fn unwrap(self) -> Vec<T> { |
||||
let err_len = self.err.len(); |
||||
if err_len > 0 { |
||||
panic!("CombinedResult has {} errors", err_len); |
||||
} |
||||
self.ok |
||||
} |
||||
|
||||
pub fn unwrap_one(self) -> T { |
||||
self.unwrap().pop().unwrap() |
||||
} |
||||
|
||||
pub fn pop_errors(&mut self) -> Vec<E> { |
||||
self.err.drain(..).collect() |
||||
} |
||||
} |
||||
|
||||
#[macro_export] |
||||
macro_rules! unwrap_enum { |
||||
($src:ident, $t:path) => { |
||||
if let $t(result) = $src { |
||||
result |
||||
} else { |
||||
panic!("wrong type") |
||||
} |
||||
}; |
||||
} |
||||
|
||||
pub fn daemonize() { |
||||
if getppid().as_raw() != 1 { |
||||
setsig(Signal::SIGTTOU, SigHandler::SigIgn); |
||||
setsig(Signal::SIGTTIN, SigHandler::SigIgn); |
||||
setsig(Signal::SIGTSTP, SigHandler::SigIgn); |
||||
} |
||||
for fd in 0..=2 { |
||||
match fdclose(fd) { |
||||
_ => (), |
||||
} |
||||
} |
||||
match chdir("/") { |
||||
_ => (), |
||||
}; |
||||
|
||||
match fork() { |
||||
Ok(ForkResult::Parent { .. }) => { |
||||
exit(0); |
||||
} |
||||
Ok(ForkResult::Child) => match setsid() { |
||||
_ => (), |
||||
}, |
||||
Err(_) => exit(255), |
||||
} |
||||
} |
||||
|
||||
pub fn setsig(sig: Signal, hnd: SigHandler) { |
||||
unsafe { |
||||
signal(sig, hnd).unwrap(); |
||||
} |
||||
} |
||||
|
||||
pub fn vec_to_string(v: &[u8]) -> String { |
||||
String::from_utf8_lossy(v).to_string() |
||||
} |
||||
|
||||
pub fn opt_to_string<T: ToString>(item: Option<T>) -> String { |
||||
match item { |
||||
Some(s) => s.to_string(), |
||||
None => String::new(), |
||||
} |
||||
} |
||||
|
||||
pub fn systime_to_string(time: &SystemTime) -> String { |
||||
DateTime::<Local>::from(*time) |
||||
.format("%d/%m/%Y %T") |
||||
.to_string() |
||||
} |
||||
|
||||
pub fn init_env() { |
||||
let envs = [".env"]; |
||||
for envfile in &envs { |
||||
dotenv::from_filename(envfile).ok(); |
||||
} |
||||
} |
||||
@ -0,0 +1,40 @@ |
||||
use crate::utils::OneOrVec; |
||||
use crate::UError; |
||||
|
||||
pub struct CombinedResult<T, E = UError> { |
||||
ok: Vec<T>, |
||||
err: Vec<E>, |
||||
} |
||||
|
||||
impl<T, E> CombinedResult<T, E> { |
||||
pub fn new() -> Self { |
||||
Self { |
||||
ok: vec![], |
||||
err: vec![], |
||||
} |
||||
} |
||||
|
||||
pub fn ok<I: OneOrVec<T>>(&mut self, result: I) { |
||||
self.ok.extend(result.into_vec()); |
||||
} |
||||
|
||||
pub fn err<I: OneOrVec<E>>(&mut self, err: I) { |
||||
self.err.extend(err.into_vec()); |
||||
} |
||||
|
||||
pub fn unwrap(self) -> Vec<T> { |
||||
let err_len = self.err.len(); |
||||
if err_len > 0 { |
||||
panic!("CombinedResult has {} errors", err_len); |
||||
} |
||||
self.ok |
||||
} |
||||
|
||||
pub fn unwrap_one(self) -> T { |
||||
self.unwrap().pop().unwrap() |
||||
} |
||||
|
||||
pub fn pop_errors(&mut self) -> Vec<E> { |
||||
self.err.drain(..).collect() |
||||
} |
||||
} |
||||
@ -0,0 +1,19 @@ |
||||
use chrono::{offset::Local, DateTime}; |
||||
use std::time::SystemTime; |
||||
|
||||
pub fn bytes_to_string(v: &[u8]) -> String { |
||||
String::from_utf8_lossy(v).to_string() |
||||
} |
||||
|
||||
pub fn opt_to_string<T: ToString>(item: Option<T>) -> String { |
||||
match item { |
||||
Some(s) => s.to_string(), |
||||
None => String::new(), |
||||
} |
||||
} |
||||
|
||||
pub fn systime_to_string(time: &SystemTime) -> String { |
||||
DateTime::<Local>::from(*time) |
||||
.format("%d/%m/%Y %T") |
||||
.to_string() |
||||
} |
||||
@ -0,0 +1,71 @@ |
||||
use nix::{ |
||||
sys::signal::{signal, SigHandler, Signal}, |
||||
unistd::{chdir, close as fdclose, fork, getppid, setsid, ForkResult}, |
||||
}; |
||||
use std::process::exit; |
||||
|
||||
pub trait OneOrVec<T> { |
||||
fn into_vec(self) -> Vec<T>; |
||||
} |
||||
|
||||
impl<T> OneOrVec<T> for T { |
||||
fn into_vec(self) -> Vec<T> { |
||||
vec![self] |
||||
} |
||||
} |
||||
|
||||
impl<T> OneOrVec<T> for Vec<T> { |
||||
fn into_vec(self) -> Vec<T> { |
||||
self |
||||
} |
||||
} |
||||
|
||||
#[macro_export] |
||||
macro_rules! unwrap_enum { |
||||
($src:ident, $t:path) => { |
||||
if let $t(result) = $src { |
||||
result |
||||
} else { |
||||
panic!("wrong type") |
||||
} |
||||
}; |
||||
} |
||||
|
||||
pub fn daemonize() { |
||||
if getppid().as_raw() != 1 { |
||||
setsig(Signal::SIGTTOU, SigHandler::SigIgn); |
||||
setsig(Signal::SIGTTIN, SigHandler::SigIgn); |
||||
setsig(Signal::SIGTSTP, SigHandler::SigIgn); |
||||
} |
||||
for fd in 0..=2 { |
||||
match fdclose(fd) { |
||||
_ => (), |
||||
} |
||||
} |
||||
match chdir("/") { |
||||
_ => (), |
||||
}; |
||||
|
||||
match fork() { |
||||
Ok(ForkResult::Parent { .. }) => { |
||||
exit(0); |
||||
} |
||||
Ok(ForkResult::Child) => match setsid() { |
||||
_ => (), |
||||
}, |
||||
Err(_) => exit(255), |
||||
} |
||||
} |
||||
|
||||
pub fn setsig(sig: Signal, hnd: SigHandler) { |
||||
unsafe { |
||||
signal(sig, hnd).unwrap(); |
||||
} |
||||
} |
||||
|
||||
pub fn init_env() { |
||||
let envs = [".env"]; |
||||
for envfile in &envs { |
||||
dotenv::from_filename(envfile).ok(); |
||||
} |
||||
} |
||||
@ -0,0 +1,11 @@ |
||||
pub mod combined_result; |
||||
pub mod conv; |
||||
pub mod misc; |
||||
pub mod tempfile; |
||||
pub mod vec_display; |
||||
|
||||
pub use combined_result::*; |
||||
pub use conv::*; |
||||
pub use misc::*; |
||||
pub use tempfile::*; |
||||
pub use vec_display::*; |
||||
@ -0,0 +1,38 @@ |
||||
use std::{env::temp_dir, fs, ops::Drop, os::unix::fs::PermissionsExt, path::PathBuf}; |
||||
use uuid::Uuid; |
||||
|
||||
pub struct TempFile { |
||||
path: PathBuf, |
||||
} |
||||
|
||||
impl TempFile { |
||||
pub fn get_path(&self) -> String { |
||||
self.path.to_string_lossy().to_string() |
||||
} |
||||
|
||||
pub fn new() -> Self { |
||||
let name = Uuid::simple(&Uuid::new_v4()).to_string(); |
||||
let mut path = temp_dir(); |
||||
path.push(name); |
||||
Self { path } |
||||
} |
||||
|
||||
pub fn write_all(&self, data: &[u8]) -> Result<(), String> { |
||||
fs::write(&self.path, data).map_err(|e| e.to_string()) |
||||
} |
||||
|
||||
pub fn write_exec(data: &[u8]) -> Result<Self, (String, String)> { |
||||
let this = Self::new(); |
||||
let path = this.get_path(); |
||||
this.write_all(data).map_err(|e| (path.clone(), e))?; |
||||
let perms = fs::Permissions::from_mode(0o555); |
||||
fs::set_permissions(&path, perms).map_err(|e| (path, e.to_string()))?; |
||||
Ok(this) |
||||
} |
||||
} |
||||
|
||||
impl Drop for TempFile { |
||||
fn drop(&mut self) { |
||||
fs::remove_file(&self.path).ok(); |
||||
} |
||||
} |
||||
@ -0,0 +1,40 @@ |
||||
use crate::{messaging::AsMsg, utils::OneOrVec}; |
||||
use serde::{Deserialize, Serialize}; |
||||
use std::fmt::{self, Display, Formatter}; |
||||
use std::ops::{Deref, DerefMut}; |
||||
|
||||
#[derive(Serialize, Deserialize, Clone, Default)] |
||||
pub struct VecDisplay<T: AsMsg + Display>(pub Vec<T>); |
||||
|
||||
impl<T: AsMsg + Display> VecDisplay<T> { |
||||
pub fn new<I: OneOrVec<T>>(inner: I) -> Self { |
||||
VecDisplay(inner.into_vec()) |
||||
} |
||||
|
||||
pub fn into_builtin_vec(self) -> Vec<T> { |
||||
self.0 |
||||
} |
||||
} |
||||
|
||||
impl<T: AsMsg + Display> Deref for VecDisplay<T> { |
||||
type Target = Vec<T>; |
||||
|
||||
fn deref(&self) -> &Self::Target { |
||||
&self.0 |
||||
} |
||||
} |
||||
|
||||
impl<T: AsMsg + Display> DerefMut for VecDisplay<T> { |
||||
fn deref_mut(&mut self) -> &mut Self::Target { |
||||
&mut self.0 |
||||
} |
||||
} |
||||
|
||||
impl<T: AsMsg + Display> Display for VecDisplay<T> { |
||||
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { |
||||
for (i, itm) in self.0.iter().enumerate() { |
||||
writeln!(f, "### {}:\n{}\n", i, itm)?; |
||||
} |
||||
Ok(()) |
||||
} |
||||
} |
||||
Loading…
Reference in new issue