Skip to content

Commit

Permalink
Improve error message in case of missing session
Browse files Browse the repository at this point in the history
Before attempting to run a command in a user session via the task
scheduler, we check if the user has an active session.
  • Loading branch information
jherbel committed Dec 5, 2023
1 parent ba6d550 commit a08e2bd
Showing 1 changed file with 71 additions and 0 deletions.
71 changes: 71 additions & 0 deletions v2/robotmk/src/bin/scheduler/sessions/schtasks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ pub fn run_task(task_spec: &TaskSpec) -> Result<RunOutcome> {
"Running the following command as task {} for user {}:\n{}\n\nBase path: {}",
task_spec.task_name, task_spec.user_name, task_spec.command_spec, task_spec.base_path
);
assert_session_is_active(task_spec.user_name)?;

let paths = Paths::from(task_spec.base_path);
create_task(task_spec, &paths)
Expand Down Expand Up @@ -89,6 +90,47 @@ impl From<&Utf8Path> for Paths {
}
}

fn assert_session_is_active(user_name: &str) -> Result<()> {
let mut query_user_command = Command::new("query");
query_user_command.arg("user");
if check_if_user_has_active_session(
user_name,
&String::from_utf8_lossy(
&query_user_command
.output()
.context(format!(
"Failed to query if user {user_name} has an active session"
))?
.stdout,
),
) {
return Ok(());
}
bail!("No active session for user {user_name} found")
}

fn check_if_user_has_active_session(user_name: &str, query_user_stdout: &str) -> bool {
for line in query_user_stdout.lines().skip(1) {
let words: Vec<&str> = line.split_whitespace().collect();
let Some(user_name_of_session) = words.first() else {
continue;
};
let Some(session_state) = words.get(3) else {
continue;
};
if (&user_name == user_name_of_session
|| &format!(
// `>` marks the current session
">{user_name}"
) == user_name_of_session)
&& session_state == &"Active"
{
return true;
}
}
false
}

fn create_task(task_spec: &TaskSpec, paths: &Paths) -> Result<()> {
write(
&paths.script,
Expand Down Expand Up @@ -254,6 +296,35 @@ mod tests {
)
}

#[test]
fn check_if_user_has_active_session_ok() {
assert!(check_if_user_has_active_session(
"vagrant",
" USERNAME SESSIONNAME ID STATE IDLE TIME LOGON TIME
>vagrant console 1 Active none 12/4/2023 9:35 AM
vagrant2 rdp-tcp#0 2 Active . 12/4/2023 9:36 AM"
))
}

#[test]
fn check_if_user_has_active_session_disconnected() {
assert!(!check_if_user_has_active_session(
"vagrant2",
" USERNAME SESSIONNAME ID STATE IDLE TIME LOGON TIME
>vagrant console 1 Active none 12/4/2023 9:35 AM
vagrant2 2 Disc . 12/4/2023 9:36 AM"
))
}

#[test]
fn check_if_user_has_active_session_no_session() {
assert!(!check_if_user_has_active_session(
"vagrant2",
" USERNAME SESSIONNAME ID STATE IDLE TIME LOGON TIME
>vagrant console 1 Active none 12/4/2023 9:35 AM"
))
}

#[test]
fn test_build_task_script() {
let mut command_spec = CommandSpec::new("C:\\somewhere\\rcc.exe");
Expand Down

0 comments on commit a08e2bd

Please sign in to comment.