bugfixes for manage-osqueryd.ps1

manage-osqueryd.ps1

#  Copyright (c) 2014-present, Facebook, Inc.
#  All rights reserved.
#
#  This source code is licensed under the BSD-style license found in the
#  LICENSE file in the root directory of this source tree. An additional grant
#  of patent rights can be found in the PATENTS file in the same directory.

param(
  [string] $args = "",
  [switch] $install = $false,
  [switch] $uninstall = $false,
  
  [switch] $start = $false,
  [switch] $stop = $false,
  
  [switch] $help = $false,
  [switch] $debug = $false
)

$kServiceName = "osquery daemon service"
$kServiceBinaryPath = Resolve-Path ([System.IO.Path]::Combine($PSScriptRoot, '..', 'osquery', 'osqueryd', 'osqueryd.exe'))
# Adapted from http://www.jonathanmedd.net/2014/01/testing-for-admin-privileges-in-powershell.html
function Test-IsAdmin {
  return ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole(
    [Security.Principal.WindowsBuiltInRole] "Administrator"
  )
}

function Do-Help {
  $programName = $MyInvocation.ScriptName#Get-Item $PSCommandPath ).Name
  
  Write-Host "Usage: $programName (-install|-uninstall|-start|-stop|-help)" -foregroundcolor Yellow
  Write-Host ""
  Write-Host "  Only one of the following options can be used. Using multiple will result in "
  Write-Host "  options being ignored."
  Write-Host "    -install          Install the osqueryd service"
  Write-Host "    -args             Specifies additional arguments for the service (only used with -install)"
  Write-Host "    -uninstall        Uninstall the osqueryd service"
  Write-Host "    -start            Start the osqueryd service"
  Write-Host "    -stop             Stop the osqueryd service"
  Write-Host ""
  Write-Host "    -help              Shows this help screen"
  
  Exit 1
}

function Do-Service {
  if (-not (Test-Path $kServiceBinaryPath)) {
    Write-Host "'$kServiceBinaryPath' is not a valid file. Did you build the osquery daemon?" -foregroundcolor Red
    Exit -1
  }
  
  $osquerydService = Get-WmiObject -Class Win32_Service -Filter "Name='$kServiceName'"
  
  if ($install) {
    if ($osquerydService) {
      Write-Host "'$kServiceName' is already installed." -foregroundcolor Yellow
      Exit 1
    } else {
      New-Service -BinaryPathName "$kServiceBinaryPath $args" -Name $kServiceName -DisplayName $kServiceName -StartupType Automatic
      Write-Host "Installed '$kServiceName' system service." -foregroundcolor Cyan
      Exit 0
    }
  } elseif ($uninstall) {
    if ($osquerydService) {
      Stop-Service $kServiceName
      
      Write-Host "Found '$kServiceName', stopping the system service..."
      
      Start-Sleep -s 5
      
      Write-Host "System service should be stopped."
      
      $osquerydService.Delete()
      Write-Host "System service '$kServiceName' uninstalled." -foregroundcolor Cyan
      
      Exit 0
    } else {
      Write-Host "'$kServiceName' is not an installed system service." -foregroundcolor Yellow
      Exit 1
    }
  } elseif ($start) {
    if ($osquerydService) {
      Start-Service $kServiceName
      Write-Host "'$kServiceName' system service is started." -foregroundcolor Cyan
    } else {
      Write-Host "'$kServiceName' is not an installed system service." -foregroundcolor Yellow
      Exit 1
    }
  } elseif ($stop) {
    if ($osquerydService) {
      Stop-Service $kServiceName
      Write-Host "'$kServiceName' system service is stopped." -foregroundcolor Cyan
    } else {
      Write-Host "'$kServiceName' is not an installed system service." -foregroundcolor Yellow
      Exit 1      
    }  
  } else {
    Write-Host "Invalid state: this should not exist!" -foregroundcolor Red
    Exit -1
  }
}

function Main {
  if (-not (Test-IsAdmin)) {
    Write-Host "Please run this script with Admin privileges!" -foregroundcolor Red
    Exit -1
  }
  
  if ($help) {
    Do-Help
  } elseif ($debug) {
    $osquerydExists = Test-Path $kServiceBinaryPath
    
    Write-Host "Service Information" -foregroundcolor Cyan
    Write-Host "  kServiceName       = '$kServiceName'" -foregroundcolor Cyan
    Write-Host "  kServiceBinaryPath = '$kServiceBinaryPath'" -foregroundcolor Cyan
    Write-Host "             +exists = $osquerydExists" -foregroundcolor Cyan
    
    Exit 0
  } elseif (([Convert]::ToInt32($install.ToBool()) + [Convert]::ToInt32($uninstall.ToBool()) + [Convert]::ToInt32($start.ToBool()) + [Convert]::ToInt32($stop.ToBool())) -Eq 1) {
    # The above is a dirty method of determining if only one of the following booleans are true.
    Do-Service
  } else {
    Write-Host "Invalid option selected: please see -help for usage details." -foregroundcolor Red
    Exit -1
  }
}

$null = Main