def download_task(url, **kwargs):
"""a simple task to use requests to get a url. By default, we return
the raw response.
Parameters
==========
REQUIRED:
url: a url to download (stream)
OPTIONAL:
write_format: to change from default "w"
disable_ssl_check: set to anything to not verify (not recommended)
"""
result = None
# Update the user what we are doing
bot.verbose("Downloading %s" % url)
# Use the basename or the user set file_name to write to
file_name = kwargs.get("file_name", os.path.basename(url))
destination = os.path.join(tempfile.gettempdir(), file_name)
verify = True
# Does the user want to disable ssl?
if "disable_ssl_check" in kwargs:
if kwargs["disable_ssl_check"]:
bot.warning(
"Verify of certificates disabled! ::TESTING USE ONLY::")
verify = False
# If the user doesn't want to write, but maybe write binary
fmt = kwargs.get("write_format", "wb")
headers = get_headers(kwargs)
# Does the url being requested exist?
if requests.head(url, verify=verify,
headers=headers).status_code in [200, 401]:
# Stream the response
response = requests.get(url,
verify=verify,
stream=True,
headers=headers)
# Invalid permissions
if response.status_code == 401:
return result
# Successful, stream to result destination
if response.status_code == 200:
chunk_size = 1 << 20
with open(destination, fmt) as filey:
for chunk in response.iter_content(chunk_size=chunk_size):
filey.write(chunk)
result = destination
return result
def has_section(self, name):
'''returns True or False to indicate if the watcher has a specified
section. To get a task, use self.has_task.
Parameters
==========
name: the name of the section to check for.
'''
self.load_config()
if name in self.config._sections:
return True
bot.warning('%s not found for watcher %s' % (name, self.name))
return False
def get_job(self, must_exist=True):
"""return the job to the user, or None"""
# Find the job based on a standard of comment
cron = self.get_crontab()
comment = "watchme-%s" % self.name
job = list(cron.find_comment(comment)) or [None]
job = job[0]
# Return None, or the actual cron job
if job is None and must_exist:
bot.warning("%s does not have a cron job configured" % self.name)
return job
def _get_pid(name):
'''used to get the pid of a process, by name
Parameters
==========
name: the name of the process to get.
'''
try:
pid = check_output(["pidof", name])
pid = pid.decode('utf-8').strip('\n').split(' ')
if len(pid) > 1:
bot.warning("More than one pid found for %s, using first." % name)
pid = int(pid[0])
except CalledProcessError:
bot.error("%s does not exist." % name)
pid = None
return pid
def finish_runs(self, results):
'''finish runs should take a dictionary of results, with keys as the
folder name, and for each, depending on the result type,
write the result to file (or update file) and then commit
to git.
Parameters
==========
results: a dictionary of tasks, with keys as the task name, and
values as the result.
'''
if results is None:
return
for name, result in results.items():
task_folder = os.path.join(self.repo, name)
if name.startswith('task'):
task = self.get_task(name)
# A decorator is run on the fly (not in config)
elif name.startswith('decorator'):
task = self.get_decorator(name)
# We only allow tasks and decorators
else:
bot.warning('%s is not task or decorator, skipping.' % name)
continue
# Ensure that the task folder exists
if not os.path.exists(task_folder):
mkdir_p(task_folder)
git_add(self.repo, task_folder)
# Files to be added to the repo via git after
files = task.write_results(result, self.repo)
# Add files to git, and commit
files.append(write_timestamp(repo=self.repo, task=name))
git_add(repo=self.repo, files=files)
git_commit(repo=self.repo,
task=self.name,
message="ADD results %s" % name)
def update_schedule(self, minute=12, hour='*', month='*', day='*'):
'''update a scheduled item from the crontab, with a new entry. This
first looks for the entry (and removes it) and then clls the new_
schedule function to write a new one. This function is intended
to be used by a client from within Python, and isn't exposed from
the command line.
'''
job = self.get_job(must_exist=True)
# Update the job if found
if job != None:
return self.schedule(minute=minute,
hour=hour,
month=month,
day=day,
job=job,
force=True)
bot.warning('%s does not have a current schedule to update.' % self.name)
def _save_file(self, result, repo, file_name=None):
"""for a result that exists, move the file to final destination.
Parameters
==========
result: the result object to save, should be path to a file
repo: the repository base with the task folder
"""
if os.path.exists(result):
name = os.path.basename(result)
file_name = self.params.get("file_name", file_name) or name
task_folder = os.path.join(repo, self.name)
# HOME/.watchme/watcher/<task>/<result>
destination = os.path.join(task_folder, file_name)
shutil.move(result, destination)
# <task>/<result>
return os.path.join(name, file_name)
bot.warning("%s does not exist." % result)
def remove_task(self, task):
'''remove a task from the watcher repo, if it exists, and the
watcher is not frozen.
Parameters
==========
task: the name of the task to remove
'''
if self.get_section(task) is not None:
if self.is_frozen():
bot.exit('watcher is frozen, unfreeze first.')
self.remove_section(task)
# If the task has a folder, remove the entire thing
repo = os.path.join(self.repo, task)
if os.path.exists(repo):
shutil.rmtree(repo)
bot.info('%s removed successfully.' % task)
git_commit(self.repo, self.name, "REMOVE task %s" % task)
else:
bot.warning('Task %s does not exist.' % task)
def monitor_pid_task(**kwargs):
'''monitor a specific process. This function can be used as a task, or
is (most likely used) for the psutils.decorators. A pid parameter
is required.
Parameters
==========
skip: an optional list of (comma separated) fields to skip. Can be in
net_io_counters,net_connections,net_if_address,net_if_stats
pid: the process id or name (required)
'''
pid = kwargs.get('pid', None)
bot.debug(kwargs)
# Helper function to get list from one,two,three
def get_list(name):
csv_list = kwargs.get(name, '')
return [x for x in csv_list.split(',') if x]
# A comma separated list of parameters to skip
skip = get_list('skip')
include = get_list('include')
only = get_list('only')
# Only continue given that argument is provided
if pid is None:
bot.warning(
"A 'pid' parameter is required to use the monitor_pid_task function."
)
return pid
# The user is allowed to provide a process name, or a number
try:
pid = int(pid)
except ValueError:
pid = _get_pid(pid)
# But if it's stil no good (None) we exit.
if pid is None:
bot.warning("'pid' must be a running process or process name.")
return pid
ps = psutil.Process(pid)
bot.debug(ps)
results = {}
for key, val in ps.as_dict().items():
# If val is None, don't include
if val is None:
bot.debug('skipping %s, None' % key)
continue
if key == "connections":
connections = []
for net in val:
entry = {
'fd': net.fd,
'family': str(net.family),
'type': str(net.type),
'laddr_ip': net.laddr.ip,
'laddr_port': net.laddr.port,
'raddr': net.raddr,
'status': net.status
}
connections.append(entry)
val = connections
# First priority goes to a custom set
if len(only) > 0:
if key in only:
results[key] = val
else:
bot.debug('skipping %s' % key)
continue
# The user requested to skip
elif key in skip or key in ['threads']:
continue
# Keep count of openfiles
elif key in ["open_files"]:
results[key] = len(val)
# Don't risk exposing sensitive information
elif key == "environ":
if key in include:
results[key] = val
# Skip over ones that are too detailed
elif key in ['memory_maps']:
continue
# I assume this included all that are in pmem too
elif isinstance(val, psutil._pslinux.pfullmem):
results[key] = {
"rss": val.rss,
"vms": val.vms,
"shared": val.shared,
"text": val.text,
"lib": val.lib,
"data": val.data,
"dirty": val.dirty,
"uss": val.uss,
"pss": val.pss,
"swap": val.swap
}
elif isinstance(val, psutil._common.pgids) or isinstance(
val, psutil._common.puids):
results[key] = {
"real": val.real,
"effetive": val.effective,
"saved": val.saved
}
elif isinstance(val, psutil._common.pcputimes):
results[key] = {
"user": val.user,
"system": val.system,
"children_user": val.children_user,
"children_system": val.children_system
}
elif isinstance(val, psutil._common.pctxsw):
results[key] = {
"voluntary": val.voluntary,
"involuntary": val.involuntary
}
elif isinstance(val, psutil._common.pionice):
results[key] = {"value": val.value}
# Older Python version (2) doesn't have attribute
if hasattr(val.ioclass, 'name'):
results[key]["ioclass"] = val.ioclass.name
# pfullmem (first above) should cover this
elif isinstance(val, psutil._pslinux.pmem):
continue
elif isinstance(val, psutil._pslinux.pio):
results[key] = {
"read_count": val.read_count,
"write_count": val.write_count,
"read_bytes": val.read_bytes,
"write_bytes": val.write_bytes,
"read_chars": val.read_chars,
"write_chars": val.write_chars
}
else:
results[key] = val
# Add any environment variables prefixed wit WATCHMEENV_
environ = get_watchme_env()
results.update(environ)
return results
