458 lines
20 KiB
Java

package name.nathanmcrae.numbersstation;
import com.leakyabstractions.result.api.Result;
import com.leakyabstractions.result.core.Results;
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.management.ManagementFactory;
import java.lang.ProcessHandle;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.time.format.DateTimeFormatter;
import java.time.LocalDateTime;
import javafx.util.Pair;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
public class WindowsScheduler {
private static final Logger logger = Logger.getLogger(Main.class.getName());
public static Document generateScheduleDocument(String authorName, String userId, StationSettings station) throws ParserConfigurationException {
DocumentBuilderFactory docFactory = DocumentBuilderFactory.newInstance();
DocumentBuilder docBuilder = docFactory.newDocumentBuilder();
// Root element
Document doc = docBuilder.newDocument();
Element rootElement = doc.createElement("Task");
rootElement.setAttribute("version", "1.2");
rootElement.setAttribute("xmlns", "http://schemas.microsoft.com/windows/2004/02/mit/task");
doc.appendChild(rootElement);
LocalDateTime scheduleDateTime = station.getScheduleStartDate().atTime(station.getScheduleStartTime());
DateTimeFormatter dateFormatter = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss");
// RegistrationInfo element
Element registrationInfo = doc.createElement("RegistrationInfo");
rootElement.appendChild(registrationInfo);
Element date = doc.createElement("Date");
date.appendChild(doc.createTextNode(scheduleDateTime.format(dateFormatter)));
registrationInfo.appendChild(date);
Element author = doc.createElement("Author");
author.appendChild(doc.createTextNode(authorName));
registrationInfo.appendChild(author);
Element uri = doc.createElement("URI");
uri.appendChild(doc.createTextNode("\\test task"));
registrationInfo.appendChild(uri);
// Triggers element
Element triggers = doc.createElement("Triggers");
rootElement.appendChild(triggers);
Element calendarTrigger = doc.createElement("CalendarTrigger");
triggers.appendChild(calendarTrigger);
Element startBoundary = doc.createElement("StartBoundary");
startBoundary.appendChild(doc.createTextNode(scheduleDateTime.format(dateFormatter)));
calendarTrigger.appendChild(startBoundary);
Element enabled = doc.createElement("Enabled");
enabled.appendChild(doc.createTextNode("true"));
calendarTrigger.appendChild(enabled);
Element randomDelay = doc.createElement("RandomDelay");
randomDelay.appendChild(doc.createTextNode("PT1H"));
calendarTrigger.appendChild(randomDelay);
switch(station.getMessagePeriod()) {
case DAILY:
{
Element scheduleByDay = doc.createElement("ScheduleByDay");
calendarTrigger.appendChild(scheduleByDay);
Element daysInterval = doc.createElement("DaysInterval");
daysInterval.appendChild(doc.createTextNode("1"));
scheduleByDay.appendChild(daysInterval);
}
break;
case WEEKLY:
{
Element scheduleByDay = doc.createElement("ScheduleByDay");
calendarTrigger.appendChild(scheduleByDay);
Element daysInterval = doc.createElement("DaysInterval");
daysInterval.appendChild(doc.createTextNode("7"));
scheduleByDay.appendChild(daysInterval);
}
break;
case MONTHLY:
Element scheduleByMonth = doc.createElement("ScheduleByMonth");
calendarTrigger.appendChild(scheduleByMonth);
Element daysOfMonth = doc.createElement("DaysOfMonth");
scheduleByMonth.appendChild(daysOfMonth);
Element day = doc.createElement("Day");
day.appendChild(doc.createTextNode(String.valueOf(station.getScheduleStartDate().getDayOfMonth())));
daysOfMonth.appendChild(day);
Element months = doc.createElement("Months");
String[] monthNames = {
"January", "February", "March", "April", "May", "June",
"July", "August", "September", "October", "November", "December"
};
for (String monthName : monthNames) {
Element month = doc.createElement(monthName);
months.appendChild(month);
}
scheduleByMonth.appendChild(months);
break;
}
// Principals element
Element principals = doc.createElement("Principals");
rootElement.appendChild(principals);
Element principal = doc.createElement("Principal");
principal.setAttribute("id", "Author");
principals.appendChild(principal);
Element userIdElement = doc.createElement("UserId");
userIdElement.appendChild(doc.createTextNode(userId));
principal.appendChild(userIdElement);
Element logonType = doc.createElement("LogonType");
logonType.appendChild(doc.createTextNode("InteractiveToken"));
principal.appendChild(logonType);
Element runLevel = doc.createElement("RunLevel");
runLevel.appendChild(doc.createTextNode("LeastPrivilege"));
principal.appendChild(runLevel);
// Settings element
Element settings = doc.createElement("Settings");
rootElement.appendChild(settings);
Element multipleInstancesPolicy = doc.createElement("MultipleInstancesPolicy");
multipleInstancesPolicy.appendChild(doc.createTextNode("IgnoreNew"));
settings.appendChild(multipleInstancesPolicy);
Element disallowStartIfOnBatteries = doc.createElement("DisallowStartIfOnBatteries");
disallowStartIfOnBatteries.appendChild(doc.createTextNode("true"));
settings.appendChild(disallowStartIfOnBatteries);
Element stopIfGoingOnBatteries = doc.createElement("StopIfGoingOnBatteries");
stopIfGoingOnBatteries.appendChild(doc.createTextNode("true"));
settings.appendChild(stopIfGoingOnBatteries);
Element allowHardTerminate = doc.createElement("AllowHardTerminate");
allowHardTerminate.appendChild(doc.createTextNode("true"));
settings.appendChild(allowHardTerminate);
Element startWhenAvailable = doc.createElement("StartWhenAvailable");
startWhenAvailable.appendChild(doc.createTextNode("false"));
settings.appendChild(startWhenAvailable);
Element runOnlyIfNetworkAvailable = doc.createElement("RunOnlyIfNetworkAvailable");
runOnlyIfNetworkAvailable.appendChild(doc.createTextNode("false"));
settings.appendChild(runOnlyIfNetworkAvailable);
Element idleSettings = doc.createElement("IdleSettings");
settings.appendChild(idleSettings);
Element stopOnIdleEnd = doc.createElement("StopOnIdleEnd");
stopOnIdleEnd.appendChild(doc.createTextNode("true"));
idleSettings.appendChild(stopOnIdleEnd);
Element restartOnIdle = doc.createElement("RestartOnIdle");
restartOnIdle.appendChild(doc.createTextNode("false"));
idleSettings.appendChild(restartOnIdle);
Element allowStartOnDemand = doc.createElement("AllowStartOnDemand");
allowStartOnDemand.appendChild(doc.createTextNode("true"));
settings.appendChild(allowStartOnDemand);
Element enabledSetting = doc.createElement("Enabled");
enabledSetting.appendChild(doc.createTextNode("true"));
settings.appendChild(enabledSetting);
Element hidden = doc.createElement("Hidden");
hidden.appendChild(doc.createTextNode("false"));
settings.appendChild(hidden);
Element runOnlyIfIdle = doc.createElement("RunOnlyIfIdle");
runOnlyIfIdle.appendChild(doc.createTextNode("false"));
settings.appendChild(runOnlyIfIdle);
Element wakeToRun = doc.createElement("WakeToRun");
wakeToRun.appendChild(doc.createTextNode("false"));
settings.appendChild(wakeToRun);
Element executionTimeLimit = doc.createElement("ExecutionTimeLimit");
executionTimeLimit.appendChild(doc.createTextNode("PT72H"));
settings.appendChild(executionTimeLimit);
Element priority = doc.createElement("Priority");
priority.appendChild(doc.createTextNode("7"));
settings.appendChild(priority);
// Actions element
Element actions = doc.createElement("Actions");
actions.setAttribute("Context", "Author");
rootElement.appendChild(actions);
Element exec = doc.createElement("Exec");
actions.appendChild(exec);
String processName = ManagementFactory.getRuntimeMXBean().getName();
long pid = Long.parseLong(processName.split("@")[0]);
ProcessHandle currentProcess = ProcessHandle.of(pid).orElseThrow();
Path executablePath = currentProcess.info().command().map(Paths::get).orElseThrow();
logger.info("Executable Path: " + executablePath);
Element command = doc.createElement("Command");
// TODO: need to figure out the real invocation
command.appendChild(doc.createTextNode(executablePath.toString()));
exec.appendChild(command);
Element arguments = doc.createElement("Arguments");
arguments.appendChild(doc.createTextNode("--station \"" + station.getName() + "\""));
exec.appendChild(arguments);
return doc;
}
public static Result<String, String> getUserId() throws IOException, InterruptedException {
Process process = new ProcessBuilder("whoami", "/user").start();
if(!process.waitFor(5, TimeUnit.SECONDS)) {
String message = "Failed to get user id: process timed out";
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
if(process.exitValue() != 0) {
Pair<String, String> output = captureProcessOutput(process);
String message = "Failed to get user id. Exit code: " + process.exitValue() + ". stdout: " + output.getKey() + "\n\tstderr: " + output.getValue();
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
try (BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()))) {
String line;
while ((line = reader.readLine()) != null) {
if (line.contains("S-")) {
return Results.success(line.split("\\s+")[1]);
}
}
}
return Results.failure("Failed to get user id: no user id found");
}
public static Result<Boolean, String> registerSchedule(StationSettings settings) {
Path tempFile = null;
try {
String taskName = "numbers-station-main_" + settings.getName();
Process queryTaskProcess = new ProcessBuilder("schtasks.exe", "/query", "/tn", taskName).start();
if (!queryTaskProcess.waitFor(5, TimeUnit.SECONDS)) {
String message = "Failed to query " + taskName + " task: process timed out";
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
// Remove previous instance of task if it exists
if (queryTaskProcess.exitValue() == 0) {
Process removeTaskProcess = new ProcessBuilder("schtasks.exe", "/delete", "/f", "/tn", taskName).start();
if (!removeTaskProcess.waitFor(5, TimeUnit.SECONDS)) {
String message = "Failed to remove " + taskName + " task: process timed out";
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
if (removeTaskProcess.exitValue() == 0) {
logger.info(taskName + " task removed successfully.");
} else {
Pair<String, String> output = captureProcessOutput(removeTaskProcess);
String message = "Failed to remove " + taskName + " task. Exit code: " + removeTaskProcess.exitValue() + ". stdout: " + output.getKey() + "\n\tstderr: " + output.getValue();
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
}
if (settings.getManageScheduleExternally()) {
return Results.success(true);
}
Result<String, String> userIdResult = getUserId();
if (!userIdResult.hasSuccess()) {
return Results.failure(userIdResult.getFailure().get());
}
String userId = userIdResult.getSuccess().get();
Document doc = generateScheduleDocument(System.getProperty("user.name"), userId, settings);
TransformerFactory transformerFactory = TransformerFactory.newInstance();
Transformer transformer = transformerFactory.newTransformer();
transformer.setOutputProperty(OutputKeys.INDENT, "yes");
transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "yes");
DOMSource source = new DOMSource(doc);
tempFile = Files.createTempFile("task", ".xml");
StreamResult result = new StreamResult(tempFile.toFile());
try (BufferedWriter writer = Files.newBufferedWriter(tempFile)) {
writer.write("<?xml version=\"1.0\"?>\n");
transformer.transform(source, new StreamResult(writer));
} catch (IOException | TransformerException e) {
String message = "Failed to write schedule to temporary file";
logger.log(Level.SEVERE, message, e);
return Results.failure(message);
}
Process process = new ProcessBuilder("schtasks.exe", "/create", "/tn", taskName, "/xml", tempFile.toString()).start();
if (!process.waitFor(5, TimeUnit.SECONDS)) {
String message = "Failed to register " + taskName + " task: process timed out";
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
if (process.exitValue() == 0) {
logger.info(taskName + " task registered successfully.");
} else {
String message = "Failed to register " + taskName + " task. Exit code: " + process.exitValue();
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
} catch (ParserConfigurationException | TransformerException | IOException | InterruptedException e) {
String message = "Exception while registering schedule";
logger.log(Level.SEVERE, message, e);
return Results.failure(message);
} finally {
if (tempFile != null) {
try {
Files.delete(tempFile);
} catch (IOException e) {
logger.log(Level.SEVERE, "Failed to delete temporary file", e);
}
}
}
return Results.success(true);
}
public static Result<Boolean, String> removeSchedule(StationSettings settings) {
Path tempFile = null;
try {
String taskName = "numbers-station-main_" + settings.getName();
Process queryTaskProcess = new ProcessBuilder("schtasks.exe", "/query", "/tn", taskName).start();
if (!queryTaskProcess.waitFor(5, TimeUnit.SECONDS)) {
String message = "Failed to query " + taskName + " task: process timed out";
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
// Remove previous instance of task if it exists
if (queryTaskProcess.exitValue() == 0) {
Process removeTaskProcess = new ProcessBuilder("schtasks.exe", "/delete", "/f", "/tn", taskName).start();
if (!removeTaskProcess.waitFor(5, TimeUnit.SECONDS)) {
String message = "Failed to remove " + taskName + " task: process timed out";
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
if (removeTaskProcess.exitValue() == 0) {
logger.info(taskName + " task removed successfully.");
} else {
Pair<String, String> output = captureProcessOutput(removeTaskProcess);
String message = "Failed to remove " + taskName + " task. Exit code: " + removeTaskProcess.exitValue() + ". stdout: " + output.getKey() + "\n\tstderr: " + output.getValue();
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
}
} catch (IOException | InterruptedException e) {
String message = "Exception while registering schedule";
logger.log(Level.SEVERE, message, e);
return Results.failure(message);
} finally {
if (tempFile != null) {
try {
Files.delete(tempFile);
} catch (IOException e) {
logger.log(Level.SEVERE, "Failed to delete temporary file", e);
}
}
}
return Results.success(true);
}
public static Result<Boolean, String> runSchedule(StationSettings settings) {
String taskName = "numbers-station-main_" + settings.getName();
try {
Process taskProcess = new ProcessBuilder("schtasks.exe", "/run", "/tn", taskName).start();
} catch (IOException e) {
String message = "Failed to run " + taskName + " task";
logger.log(Level.SEVERE, message, e);
return Results.failure(message);
}
return Results.success(true);
}
public static Result<Boolean, String> scheduleExists(StationSettings settings) {
try {
String taskName = "numbers-station-main_" + settings.getName();
Process queryTaskProcess = new ProcessBuilder("schtasks.exe", "/query", "/tn", taskName).start();
if (!queryTaskProcess.waitFor(5, TimeUnit.SECONDS)) {
String message = "Failed to query " + taskName + " task: process timed out";
logger.log(Level.SEVERE, message);
return Results.failure(message);
}
return Results.success(queryTaskProcess.exitValue() == 0);
} catch (InterruptedException | IOException e) {
String message = "Exception while querying schedule";
logger.log(Level.SEVERE, message, e);
return Results.failure(message);
}
}
/**
* @return (stdout contents, stderr contents)
* TODO: don't assume that process has exited yet. If it does we don't want to hang.
*/
public static Pair<String, String> captureProcessOutput(Process process) throws IOException {
StringBuilder output = new StringBuilder();
StringBuilder error = new StringBuilder();
try (BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
BufferedReader errorReader = new BufferedReader(new InputStreamReader(process.getErrorStream()))) {
String line;
while ((line = reader.readLine()) != null) {
output.append(line + "\n");
}
while ((line = errorReader.readLine()) != null) {
error.append(line + "\n");
}
}
return new Pair<>(output.toString(), error.toString());
}
}