metasploit-framework/lib/rex/job_container.rb

201 lines
4.1 KiB
Ruby

# -*- coding: binary -*-
module Rex
###
#
# This class is the concrete representation of an abstract job.
#
###
class Job
#
# Creates an individual job instance and initializes it with the supplied
# parameters.
#
def initialize(container, jid, name, ctx, run_proc, clean_proc)
self.container = container
self.jid = jid
self.name = name
self.run_proc = run_proc
self.clean_proc = clean_proc
self.ctx = ctx
self.start_time = nil
end
#
# Runs the job in the context of its own thread if the async flag is false.
# Otherwise, the job is run inline.
#
def start(async = false)
self.start_time = Time.now
if (async)
self.job_thread = Rex::ThreadFactory.spawn("JobID(#{jid})-#{name}", false) {
# Deschedule our thread momentarily
::IO.select(nil, nil, nil, 0.01)
begin
run_proc.call(ctx)
ensure
clean_proc.call(ctx)
container.remove_job(self)
end
}
else
begin
run_proc.call(ctx)
rescue ::Exception
container.stop_job(jid)
raise $!
end
end
end
#
# Stops the job if it's currently running and calls its cleanup procedure
#
def stop
if (self.job_thread)
self.job_thread.kill
self.job_thread = nil
end
clean_proc.call(ctx) if (clean_proc)
end
#
# The name of the job.
#
attr_reader :name
#
# The job identifier as assigned by the job container.
#
attr_reader :jid
#
# The time at which this job was started.
#
attr_reader :start_time
#
# Some job context.
#
attr_reader :ctx
protected
attr_writer :name #:nodoc:
attr_writer :jid #:nodoc:
attr_accessor :job_thread #:nodoc:
attr_accessor :container #:nodoc:
attr_accessor :run_proc #:nodoc:
attr_accessor :clean_proc #:nodoc:
attr_writer :ctx #:nodoc:
attr_writer :start_time #:nodoc:
end
###
#
# This class contains zero or more abstract jobs that can be enumerated and
# stopped in a generic fashion. This is used to provide a mechanism for
# keeping track of arbitrary contexts that may or may not require a dedicated
# thread.
#
###
class JobContainer < Hash
def initialize
self.job_id_pool = 0
end
#
# Adds an already running task as a symbolic job to the container.
#
def add_job(name, ctx, run_proc, clean_proc)
real_name = name
count = 0
jid = job_id_pool
self.job_id_pool += 1
# If we were not supplied with a job name, pick one from the hat
if (real_name == nil)
real_name = '#' + jid.to_s
end
# Find a unique job name
while (j = self[real_name])
real_name = name + " #{count}"
count += 1
end
j = Job.new(self, jid, real_name, ctx, run_proc, clean_proc)
self[jid.to_s] = j
end
#
# Starts a job using the supplied name and run/clean procedures.
#
def start_job(name, ctx, run_proc, clean_proc = nil)
j = add_job(name, ctx, run_proc, clean_proc)
j.start
j.jid
end
#
# Starts a background job that doesn't call the cleanup routine or run
# the run_proc in its own thread. Rather, the run_proc is called
# immediately and the clean_proc is never called until the job is removed
# from the job container.
#
def start_bg_job(name, ctx, run_proc, clean_proc = nil, async = true)
j = add_job(name, ctx, run_proc, clean_proc)
j.start(async)
j.jid
end
#
# Stops the job with the supplied name and forces it to cleanup. Stopping
# the job also leads to its removal.
#
def stop_job(jid)
if (j = self[jid.to_s])
j.stop
remove_job(j)
end
end
#
# Removes a job that was previously running. This is typically called when
# a job completes its task.
#
def remove_job(inst)
self.delete(inst.jid.to_s)
end
#
# Overrides the builtin 'each' operator to avoid the following exception on Ruby 1.9.2+
# "can't add a new key into hash during iteration"
#
def each(&block)
list = []
self.keys.sort.each do |sidx|
list << [sidx, self[sidx]]
end
list.each(&block)
end
protected
attr_accessor :job_id_pool # :nodoc:
end
end