Package org.openqa.selenium.browserlaunchers

Source Code of org.openqa.selenium.browserlaunchers.WindowsUtils

/*
* Copyright 2006 ThoughtWorks, Inc.
*
*  Licensed under the Apache License, Version 2.0 (the "License");
*  you may not use this file except in compliance with the License.
*  You may obtain a copy of the License at
*
*      http://www.apache.org/licenses/LICENSE-2.0
*
*  Unless required by applicable law or agreed to in writing, software
*  distributed under the License is distributed on an "AS IS" BASIS,
*  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*  See the License for the specific language governing permissions and
*  limitations under the License.
*
*/
package org.openqa.selenium.browserlaunchers;

import javax.xml.parsers.DocumentBuilderFactory;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.openqa.selenium.internal.CommandLine;
import org.openqa.selenium.internal.Trace;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;
import org.w3c.dom.Text;

public class WindowsUtils {

  public static Boolean regVersion1 = null;
 
  private static Trace log = new NullTrace();
  private static final boolean THIS_IS_WINDOWS = File.pathSeparator.equals(";");
  private static String wmic = null;
  private static File wbem = null;
  private static String taskkill = null;
  private static String reg = null;
  private static Properties env = null;

  /**
   * @param args
   * @throws IOException
   */
  public static void main(String[] args) throws Exception {
    if (args.length == 0) {
      System.out.println("Kills Windows processes by matching their command lines");
      System.out.println("usage: " + WindowsUtils.class.getName() + " command arg1 arg2 ...");
    }
    kill(args);

  }

  public static void traceWith(Trace log) {
    WindowsUtils.log = log;
  }

  /**
   * Kill processes by name
   */
  public static void killByName(String name) {
    executeCommand("taskkill", "/f", "/im", name);
  }

  /**
   * Kill processes by name, log and ignore errors
   */
  public static void tryToKillByName(String name) {
    if (!thisIsWindows()) {
      return;
    }
    try {
      killByName(name);
    } catch (WindowsRegistryException e) {
      log.warn(e);
    }
  }

  /**
   * Searches the process list for a process with the specified command line and kills it
   *
   * @param cmdarray the array of command line arguments
   * @throws Exception if something goes wrong while reading the process list or searching for your command line
   */
  public static void kill(String[] cmdarray) throws Exception {
    StringBuffer pattern = new StringBuffer();
    File executable = new File(cmdarray[0]);
    /* For the first argument, the executable, Windows may modify
    * the start path in any number of ways.  Ignore a starting quote
    * if any (\"?), non-greedily look for anything up until the last
    * backslash (.*?\\\\), then look for the executable's filename,
    * then finally ignore a final quote (\"?)
    */
    // TODO We should be careful, in case Windows has ~1-ified the executable name as well
    pattern.append("\"?.*?\\\\");
    pattern.append(executable.getName());
    pattern.append("\"?");
    for (int i = 1; i < cmdarray.length; i++) {
      /* There may be a space, but maybe not (\\s?), may be a quote or maybe not (\"?),
      * but then turn on block quotation (as if *everything* had a regex backslash in front of it)
      * with \Q.  Then look for the next argument (which may have ?s, \s, "s, who knows),
      * turning off block quotation.  Now ignore a final quote if any (\"?)
      */
      pattern.append("\\s?\"?\\Q");
      String arg = cmdarray[i];
      pattern.append(arg);
      pattern.append("\\E\"?");
    }
    pattern.append("\\s*");
    Pattern cmd = Pattern.compile(pattern.toString(), Pattern.CASE_INSENSITIVE);
    Map procMap = procMap();
    boolean killedOne = false;
    for (Iterator i = procMap.keySet().iterator(); i.hasNext();) {
      String commandLine = (String) i.next();
      if (commandLine == null) {
        continue;
      }
      Matcher m = cmd.matcher(commandLine);
      if (m.matches()) {
        String processID = (String) procMap.get(commandLine);
        StringBuilder logMessage = new StringBuilder("Killing PID ");
        logMessage.append(processID);
        logMessage.append(": ");
        logMessage.append(commandLine);
        log.info(logMessage.toString());
        killPID(processID);
        log.info("Killed");
        killedOne = true;
      }
    }
    if (!killedOne) {
      StringBuilder errorMessage = new StringBuilder("Didn't find any matches for");
      for (int i = 0; i < cmdarray.length; i++) {
        errorMessage.append(" '");
        errorMessage.append(cmdarray[i]);
        errorMessage.append('\'');
      }
      log.warn(errorMessage.toString());
    }
  }

  /**
   * Kills the specified process ID
   */
  private static void killPID(String processID) {
    executeCommand("taskkill", "/pid", processID);
  }

  /**
   * Returns a map of process IDs to command lines
   *
   * @return a map of process IDs to command lines
   * @throws Exception - if something goes wrong while reading the process list
   */
  public static Map procMap() throws Exception {
    log.info("Reading Windows Process List...");
    String output = executeCommand(findWMIC(), "process", "list", "full", "/format:rawxml.xsl");
//    exec.setFailonerror(true);
    log.info("Done, searching for processes to kill...");
    // WMIC drops an ugly zero-length batch file; clean that up
    File TempWmicBatchFile = new File("TempWmicBatchFile.bat");
    if (TempWmicBatchFile.exists()) {
      TempWmicBatchFile.delete();
    }

    // TODO This would be faster if it used SAX instead of DOM
    Document doc = DocumentBuilderFactory.newInstance().newDocumentBuilder()
        .parse(new ByteArrayInputStream(output.getBytes()));
    NodeList procList = doc.getElementsByTagName("INSTANCE");
    Map<String, String> processes = new HashMap<String, String>();
    for (int i = 0; i < procList.getLength(); i++) {
      Element process = (Element) procList.item(i);
      NodeList propList = process.getElementsByTagName("PROPERTY");
      Map<String, Object> procProps = new HashMap<String, Object>();
      for (int j = 0; j < propList.getLength(); j++) {
        Element property = (Element) propList.item(j);
        String propName = property.getAttribute("NAME");
        NodeList valList = property.getElementsByTagName("VALUE");
        String value = null;
        if (valList.getLength() != 0) {
          Element valueElement = (Element) valList.item(0);
          Text valueNode = (Text) valueElement.getFirstChild();
          value = valueNode.getData();
        }
        procProps.put(propName, value);
      }
      String processID = (String) procProps.get("ProcessId");
      String commandLine = (String) procProps.get("CommandLine");
      processes.put(commandLine, processID);
    }
    return processes;
  }

  /**
   * Returns the current process environment variables
   *
   * @return the current process environment variables
   */
  public static synchronized Properties loadEnvironment() {
    if (env != null) {
      return env;
    }
    env = new Properties();
    for (Map.Entry<String, String> entry : System.getenv().entrySet()) {
      env.put(entry.getKey(), entry.getValue());
    }
    return env;
  }

  /**
   * Retrieve the exact case-sensitive name of the "Path" environment variable,
   * which may be any one of "PATH", "Path" or "path".
   *
   * @return the exact case-sensitive name of the "Path" environment variable
   */
  public static String getExactPathEnvKey() {
    loadEnvironment();
    for (Iterator i = env.keySet().iterator(); i.hasNext();) {
      String key = (String) i.next();
      if (key.equalsIgnoreCase("PATH")) {
        return key;
      }
    }
    // They don't have a path???
    return "PATH";
  }

  public static String getPath() {
    loadEnvironment();
    return getEnvVarIgnoreCase("PATH");
  }

  /**
   * Returns the path to the Windows Program Files.  On non-English versions,
   * this is not necessarily "C:\Program Files".
   *
   * @return the path to the Windows Program Files
   */
  public static String getProgramFilesPath() {
    loadEnvironment();
    String pf = getEnvVarIgnoreCase("ProgramFiles");
    if (pf != null) {
      File ProgramFiles = new File(pf);
      if (ProgramFiles.exists()) {
        return ProgramFiles.getAbsolutePath();
      }
    }
    return new File("C:\\Program Files").getAbsolutePath();
  }

  /**
   * Returns the path to Local AppData.  For different users, this will be
   * different.
   *
   * @return the path to Local AppData
   */
  public static String getLocalAppDataPath() {
    loadEnvironment();
    final String keyLocalAppData =
        "HKEY_CURRENT_USER\\Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\Shell Folders\\Local AppData";
    String localAppDataPath = readStringRegistryValue(keyLocalAppData);
    String userProfile = getEnvVarIgnoreCase("USERPROFILE");
    if (userProfile != null) {
      return localAppDataPath.replace("%USERPROFILE%", userProfile);
    }
    return localAppDataPath;
  }

  public static String getEnvVarIgnoreCase(String var) {
    loadEnvironment();
    for (Iterator i = env.keySet().iterator(); i.hasNext();) {
      String key = (String) i.next();
      if (key.equalsIgnoreCase(var)) {
        return env.getProperty(key);
      }
    }
    return null;
  }

  /**
   * Finds the system root directory, e.g. "c:\windows" or "c:\winnt"
   */
  public static File findSystemRoot() {
    Properties p = loadEnvironment();
    String systemRootPath = (String) p.get("SystemRoot");
    if (systemRootPath == null) {
      systemRootPath = (String) p.get("SYSTEMROOT");
    }
    if (systemRootPath == null) {
      systemRootPath = (String) p.get("systemroot");
    }
    if (systemRootPath == null) {
      throw new RuntimeException("SystemRoot apparently not set!");
    }
    File systemRoot = new File(systemRootPath);
    if (!systemRoot.exists()) {
      throw new RuntimeException("SystemRoot doesn't exist: " + systemRootPath);
    }
    return systemRoot;
  }

  /**
   * Finds WMIC.exe
   *
   * @return the exact path to wmic.exe, or just the string "wmic" if it couldn't be found (in which case you can pass that to exec to try to run it from the path)
   */
  public static String findWMIC() {
    if (wmic != null) {
      return wmic;
    }
    findWBEM();
    if (null != wbem) {
      File wmicExe = new File(findWBEM(), "wmic.exe");
      if (wmicExe.exists()) {
        wmic = wmicExe.getAbsolutePath();
        return wmic;
      }
    }
    log.warn("Couldn't find wmic! Hope it's on the path...");
    wmic = "wmic";
    return wmic;
  }

  /**
   * Finds the WBEM directory in the systemRoot directory
   *
   * @return the WBEM directory, or <code>null</code> if it couldn't be found
   */
  public static File findWBEM() {
    if (wbem != null) {
      return wbem;
    }
    File systemRoot = findSystemRoot();
    wbem = new File(systemRoot, "system32/wbem");
    if (!wbem.exists()) {
      log.error("Couldn't find wbem!");
      return null;
    }
    return wbem;
  }

  /**
   * Finds taskkill.exe
   *
   * @return the exact path to taskkill.exe, or just the string "taskkill" if it couldn't be found (in which case you can pass that to exec to try to run it from the path)
   */
  public static String findTaskKill() {
    if (taskkill != null) {
      return taskkill;
    }
    File systemRoot = findSystemRoot();
    File taskkillExe = new File(systemRoot, "system32/taskkill.exe");
    if (taskkillExe.exists()) {
      taskkill = taskkillExe.getAbsolutePath();
      return taskkill;
    }
    log.warn("Couldn't find taskkill! Hope it's on the path...");
    taskkill = "taskkill";
    return taskkill;
  }

  /**
   * Finds reg.exe
   *
   * @return the exact path to reg.exe, or just the string "reg" if it couldn't be found (in which case you can pass that to exec to try to run it from the path)
   */
  public static String findReg() {
    if (reg != null) {
      return reg;
    }
    File systemRoot = findSystemRoot();
    File regExe = new File(systemRoot, "system32/reg.exe");
    if (regExe.exists()) {
      reg = regExe.getAbsolutePath();
      return reg;
    }
    regExe = new File("c:\\ntreskit\\reg.exe");
    if (regExe.exists()) {
      reg = regExe.getAbsolutePath();
      return reg;
    }
    reg = CommandLine.findExecutable("reg.exe");
    if (reg != null) {
      return reg;
    }
    log.error("OS Version: " + System.getProperty("os.version"));
    throw new WindowsRegistryException("Couldn't find reg.exe!\n" +
                                       "Please download it from Microsoft and install it in a standard location.\n"
                                       +
                                       "See here for details: http://wiki.openqa.org/display/SRC/Windows+Registry+Support");
  }

  public static boolean isRegExeVersion1() {
    if (regVersion1 != null) {
      return regVersion1.booleanValue();
    }

    String output = executeCommand(findReg(), "/?");
    boolean version1 = output.indexOf("version 1.0") != -1;
    regVersion1 = new Boolean(version1);
    return version1;
  }

  public static Class discoverRegistryKeyType(String key) {
    if (!doesRegistryValueExist(key)) {
      return null;
    }
    RegKeyValue r = new RegKeyValue(key);
    String output = runRegQuery(key);
    Pattern pat;
    if (isRegExeVersion1()) {
      pat = Pattern.compile("\\s*(REG_\\S+)");
    } else {
      pat = Pattern.compile("\\Q" + r.value + "\\E\\s+(REG_\\S+)\\s+(.*)");
    }
    Matcher m = pat.matcher(output);
    if (!m.find()) {
      throw new WindowsRegistryException("Output didn't look right: " + output);
    }
    String type = m.group(1);
    if ("REG_SZ".equals(type) || "REG_EXPAND_SZ".equals(type)) {
      return String.class;
    } else if ("REG_DWORD".equals(type)) {
      return int.class;
    } else {
      throw new WindowsRegistryException("Unknown type: " + type);
    }
  }

  public static String readStringRegistryValue(String key) {
    RegKeyValue r = new RegKeyValue(key);
    String output = runRegQuery(key);
    Pattern pat;
    if (isRegExeVersion1()) {
      pat = Pattern.compile("\\s*(REG_\\S+)\\s+\\Q" + r.value + "\\E\\s+(.*)");
    } else {
      pat = Pattern.compile("\\Q" + r.value + "\\E\\s+(REG_\\S+)\\s+(.*)");
    }
    Matcher m = pat.matcher(output);
    if (!m.find()) {
      throw new WindowsRegistryException("Output didn't look right: " + output);
    }
    String type = m.group(1);
    if (!"REG_SZ".equals(type) && !"REG_EXPAND_SZ".equals(type)) {
      throw new WindowsRegistryException(
          r.value + " was not a REG_SZ or a REG_EXPAND_SZ (String): " + type);
    }
    String value = m.group(2);
    return value;
  }

  public static int readIntRegistryValue(String key) {
    RegKeyValue r = new RegKeyValue(key);
    String output = runRegQuery(key);
    Pattern pat;
    if (isRegExeVersion1()) {
      pat = Pattern.compile("\\s*(REG_\\S+)\\s+\\Q" + r.value + "\\E\\s+(.*)");
    } else {
      pat = Pattern.compile("\\Q" + r.value + "\\E\\s+(REG_\\S+)\\s+0x(.*)");
    }

    Matcher m = pat.matcher(output);
    if (!m.find()) {
      throw new WindowsRegistryException("Output didn't look right: " + output);
    }
    String type = m.group(1);
    if (!"REG_DWORD".equals(type)) {
      throw new WindowsRegistryException(r.value + " was not a REG_DWORD (int): " + type);
    }
    String strValue = m.group(2);
    int value;
    if (isRegExeVersion1()) {
      value = Integer.parseInt(strValue);
    } else {
      value = Integer.parseInt(strValue, 16);
    }
    return value;
  }

  public static boolean readBooleanRegistryValue(String key) {
    RegKeyValue r = new RegKeyValue(key);
    int value = readIntRegistryValue(key);
    if (0 == value) {
      return false;
    }
    if (1 == value) {
      return true;
    }
    throw new WindowsRegistryException(r.value + " was not either 0 or 1: " + value);
  }

  public static boolean doesRegistryValueExist(String key) {
    List<String> args = new ArrayList();
    args.add("query");

    if (isRegExeVersion1()) {
      args.add(key);
    } else {
      RegKeyValue r = new RegKeyValue(key);
      args.add(r.key);
      args.add("/v");
      args.add(r.value);
    }

    try {
      executeCommand(findReg(), args.toArray(new String[args.size()]));
      return true;
    } catch (WindowsRegistryException e) {
      return false;
    }
  }

  public static void writeStringRegistryValue(String key, String data)
      throws WindowsRegistryException {
    List<String> args = new ArrayList<String>();
    if (isRegExeVersion1()) {
      if (doesRegistryValueExist(key)) {
        args.add("update");
      } else {
        args.add("add");
      }
      args.add(key + "=" + data);
    } else {
      args.add("add");
      RegKeyValue r = new RegKeyValue(key);
      args.add(r.key);
      args.add("/v");
      args.add(r.value);
      args.add("/d");
      args.add(data);
      args.add("/f");
    }

    executeCommand(findReg(), args.toArray(new String[args.size()]));
  }

  private static String executeCommand(String commandName, String... args) {
    CommandLine cmd = new CommandLine(commandName, args);
    cmd.execute();

    String output = cmd.getStdOut();
    if (!cmd.isSuccessful()) {
      throw new WindowsRegistryException("exec return code " + cmd.getExitCode() + ": " + output);
    }
    return output;
  }

  public static void writeIntRegistryValue(String key, int data) {
    List<String> args = new ArrayList<String>();
    if (isRegExeVersion1()) {
      if (doesRegistryValueExist(key)) {
        args.add("update");
        args.add(key + "=" + Integer.toString(data));
      } else {
        args.add("add");
        args.add(key + "=" + Integer.toString(data));
        args.add("REG_DWORD");
      }
    } else {
      args.add("add");
      RegKeyValue r = new RegKeyValue(key);
      args.add(r.key);
      args.add("/v");
      args.add(r.value);
      args.add("/t");
      args.add("REG_DWORD");
      args.add("/d");
      args.add(Integer.toString(data));
      args.add("/f");
    }

    executeCommand(findReg(), args.toArray(new String[args.size()]));
  }

  public static void writeBooleanRegistryValue(String key, boolean data) {
    writeIntRegistryValue(key, data ? 1 : 0);
  }

  public static void deleteRegistryValue(String key) {
    List<String> args = new ArrayList<String>();
    if (isRegExeVersion1()) {
      args.add("delete");
      args.add(key);
      args.add("/FORCE");
    } else {
      RegKeyValue r = new RegKeyValue(key);
      args.add("delete");
      args.add(r.key);
      args.add("/v");
      args.add(r.value);
      args.add("/f");
    }

    executeCommand(findReg(), args.toArray(new String[args.size()]));
  }

  /**
   * Executes reg.exe to query the registry
   */
  private static String runRegQuery(String key) {
    List<String> args = new ArrayList<String>();
    args.add("query");
    if (isRegExeVersion1()) {
      args.add(key);
    } else {
      RegKeyValue r = new RegKeyValue(key);
      args.add(r.key);
      args.add("/v");
      args.add(r.value);
    }

    return executeCommand(findReg(), args.toArray(new String[args.size()]));
  }

  private static class RegKeyValue {
    private String key;
    private String value;

    public RegKeyValue(String path) {
      int i = path.lastIndexOf('\\');
      key = path.substring(0, i);
      value = path.substring(i + 1);
    }
  }

  /**
   * Returns true if the current OS is MS Windows; false otherwise
   *
   * @return true if the current OS is MS Windows; false otherwise
   */
  public static boolean thisIsWindows() {
    return THIS_IS_WINDOWS;
  }

}
TOP

Related Classes of org.openqa.selenium.browserlaunchers.WindowsUtils

TOP
Copyright © 2018 www.massapi.com. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.