nutshell 1.0.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,440 @@
1
+ module Nutshell
2
+
3
+ ##
4
+ # The Shell class handles local input, output and execution to the shell.
5
+
6
+ class Shell
7
+
8
+ include Open4
9
+
10
+ LOCAL_USER = `whoami`.chomp
11
+ LOCAL_HOST = `hostname`.chomp
12
+
13
+ SUDO_FAILED = /^Sorry, try again./
14
+ SUDO_PROMPT = /^Password:/
15
+
16
+ attr_reader :user, :host, :password, :input, :output, :mutex
17
+ attr_accessor :env, :sudo, :timeout
18
+
19
+ def initialize output=$stdout, options={}
20
+ @output = output
21
+
22
+ $stdin.sync
23
+ @input = HighLine.new $stdin
24
+
25
+ @user = LOCAL_USER
26
+ @host = LOCAL_HOST
27
+
28
+ @sudo = options[:sudo]
29
+ @env = options[:env] || {}
30
+ @password = options[:password]
31
+
32
+ @timeout = options[:timeout] || Nutshell.timeout
33
+
34
+ @cmd_activity = nil
35
+
36
+ @mutex = nil
37
+ end
38
+
39
+
40
+ ##
41
+ # Checks for equality
42
+
43
+ def == shell
44
+ @host == shell.host && @user == shell.user rescue false
45
+ end
46
+
47
+
48
+ ##
49
+ # Prompt the user for input.
50
+
51
+ def ask(*args, &block)
52
+ sync{ @input.ask(*args, &block) }
53
+ end
54
+
55
+
56
+ ##
57
+ # Prompt the user to agree.
58
+
59
+ def agree(*args, &block)
60
+ sync{ @input.agree(*args, &block) }
61
+ end
62
+
63
+
64
+ ##
65
+ # Execute a command on the local system and return the output.
66
+
67
+ def call cmd, options={}, &block
68
+ execute sudo_cmd(cmd, options), &block
69
+ end
70
+
71
+
72
+ ##
73
+ # Close the output IO. (Required by the Logger class)
74
+
75
+ def close
76
+ @output.close
77
+ end
78
+
79
+
80
+ ##
81
+ # Returns true. Compatibility method with RemoteShell.
82
+
83
+ def connect
84
+ true
85
+ end
86
+
87
+
88
+ ##
89
+ # Returns true. Compatibility method with RemoteShell.
90
+
91
+ def connected?
92
+ true
93
+ end
94
+
95
+
96
+ ##
97
+ # Returns true. Compatibility method with RemoteShell.
98
+
99
+ def disconnect
100
+ true
101
+ end
102
+
103
+
104
+ ##
105
+ # Copies a file. Compatibility method with RemoteShell.
106
+
107
+ def download from_path, to_path, options={}, &block
108
+ FileUtils.cp_r from_path, to_path
109
+ end
110
+
111
+ alias upload download
112
+
113
+
114
+ ##
115
+ # Expands the path. Compatibility method with RemoteShell.
116
+
117
+ def expand_path path
118
+ File.expand_path path
119
+ end
120
+
121
+
122
+ ##
123
+ # Checks if file exists. Compatibility method with RemoteShell.
124
+
125
+ def file? filepath
126
+ File.file? filepath
127
+ end
128
+
129
+
130
+ ##
131
+ # Start an interactive shell with preset permissions and env.
132
+ # Optionally pass a command to be run first.
133
+
134
+ def tty! cmd=nil
135
+ sync do
136
+ cmd = [cmd, "sh -il"].compact.join " && "
137
+ pid = fork do
138
+ exec sudo_cmd(env_cmd(cmd)).to_a.join(" ")
139
+ end
140
+ Process.waitpid pid
141
+ end
142
+ end
143
+
144
+
145
+ ##
146
+ # Write a file. Compatibility method with RemoteShell.
147
+
148
+ def make_file filepath, content, options={}
149
+ File.open(filepath, "w+"){|f| f.write(content)}
150
+ end
151
+
152
+
153
+ ##
154
+ # Get the name of the OS
155
+
156
+ def os_name
157
+ @os_name ||= call("uname -s").strip.downcase
158
+ end
159
+
160
+
161
+ ##
162
+ # Prompt the user for a password
163
+
164
+ def prompt_for_password
165
+ host_info = [@user, @host].compact.join("@")
166
+ @password = ask("#{host_info} Password:") do |q|
167
+ q.echo = false
168
+ end
169
+ end
170
+
171
+
172
+ ##
173
+ # Build an env command if an env_hash is passed
174
+
175
+ def env_cmd cmd, env_hash=@env
176
+ if env_hash && !env_hash.empty?
177
+ env_vars = env_hash.map{|e| e.join("=")}
178
+ cmd = ["env", env_vars, cmd].flatten
179
+ end
180
+ cmd
181
+ end
182
+
183
+
184
+ ##
185
+ # Wrap command in quotes and escape as needed.
186
+
187
+ def quote_cmd cmd
188
+ cmd = [*cmd].join(" ")
189
+ "'#{cmd.gsub(/'/){|s| "'\\''"}}'"
190
+ end
191
+
192
+
193
+ ##
194
+ # Runs the given block within a session.
195
+ # Will not disconnect if previous session had been started.
196
+
197
+ def session &block
198
+ was_connected = connected?
199
+ connect
200
+ yield self if block_given?
201
+ disconnect unless was_connected
202
+ end
203
+
204
+
205
+ ##
206
+ # Build an sh -c command
207
+
208
+ def sh_cmd cmd
209
+ ["sh", "-c", quote_cmd(cmd)]
210
+ end
211
+
212
+
213
+ ##
214
+ # Build a command with sudo.
215
+ # If sudo_val is nil, it is considered to mean "pass-through"
216
+ # and the default shell sudo will be used.
217
+ # If sudo_val is false, the cmd will be returned unchanged.
218
+ # If sudo_val is true, the returned command will be prefaced
219
+ # with sudo -H
220
+ # If sudo_val is a String, the command will be prefaced
221
+ # with sudo -H -u string_value
222
+
223
+ def sudo_cmd cmd, sudo_val=nil
224
+ sudo_val = sudo_val[:sudo] if Hash === sudo_val
225
+ sudo_val = @sudo if sudo_val.nil?
226
+
227
+ case sudo_val
228
+ when true
229
+ ["sudo", "-H", cmd].flatten
230
+ when String
231
+ ["sudo", "-H", "-u", sudo_val, cmd].flatten
232
+ else
233
+ cmd
234
+ end
235
+ end
236
+
237
+
238
+ ##
239
+ # Force symlinking a directory.
240
+
241
+ def symlink target, symlink_name
242
+ call "ln -sfT #{target} #{symlink_name}" rescue false
243
+ end
244
+
245
+
246
+ ##
247
+ # Synchronize a block with the current mutex if it exists.
248
+
249
+ def sync
250
+ if @mutex
251
+ @mutex.synchronize{ yield }
252
+ else
253
+ yield
254
+ end
255
+ end
256
+
257
+
258
+ ##
259
+ # Returns true if command was run successfully, otherwise returns false.
260
+
261
+ def syscall cmd, options=nil
262
+ call(cmd, options) && true rescue false
263
+ end
264
+
265
+
266
+ ##
267
+ # Checks if timeout occurred.
268
+
269
+ def timed_out? start_time=@cmd_activity, max_time=@timeout
270
+ return unless max_time
271
+ Time.now.to_i - start_time.to_i > max_time
272
+ end
273
+
274
+
275
+ ##
276
+ # Update the time of the last command activity
277
+
278
+ def update_timeout
279
+ @cmd_activity = Time.now
280
+ end
281
+
282
+
283
+ ##
284
+ # Execute a block while setting the shell's mutex.
285
+ # Sets the mutex to its original value on exit.
286
+ # Executing commands with a mutex is used for user prompts.
287
+
288
+ def with_mutex mutex
289
+ old_mutex, @mutex = @mutex, mutex
290
+ yield
291
+ @mutex = old_mutex
292
+ end
293
+
294
+
295
+ ##
296
+ # Runs the passed block within a connection session.
297
+ # If the shell is already connected, connecting and disconnecting
298
+ # is ignored; otherwise, the session method will ensure that
299
+ # the shell's connection gets closed after the block has been
300
+ # executed.
301
+
302
+ def with_session
303
+ prev_connection = connected?
304
+ connect unless prev_connection
305
+
306
+ yield
307
+
308
+ disconnect unless prev_connection
309
+ end
310
+
311
+
312
+ ##
313
+ # Write string to stdout (by default).
314
+
315
+ def write str
316
+ @output.write str
317
+ end
318
+
319
+ alias << write
320
+
321
+
322
+ ##
323
+ # Execute a command with open4 and loop until the process exits.
324
+ # The cmd argument may be a string or an array. If a block is passed,
325
+ # it will be called when data is received and passed the stream type
326
+ # and stream string value:
327
+ # shell.execute "test -s 'blah' && echo 'true'" do |stream, str|
328
+ # stream #=> :stdout
329
+ # string #=> 'true'
330
+ # end
331
+ #
332
+ # The method returns the output from the stdout stream by default, and
333
+ # raises a CmdError if the exit status of the command is not zero.
334
+
335
+ def execute cmd
336
+ cmd = [cmd] unless Array === cmd
337
+ pid, inn, out, err = popen4(*cmd)
338
+
339
+ inn.sync = true
340
+ log_methods = {out => :debug, err => :error}
341
+
342
+ result, status = process_streams(pid, out, err) do |stream, data|
343
+ stream_name = :out if stream == out
344
+ stream_name = :err if stream == err
345
+ stream_name = :inn if stream == inn
346
+
347
+
348
+ # User blocks should run with sync threads to avoid badness.
349
+ sync do
350
+ yield(stream_name, data, inn) if block_given?
351
+ end
352
+
353
+
354
+ if password_required?(stream_name, data) then
355
+
356
+ kill_process(pid) unless Nutshell.interactive?
357
+
358
+ send_password_to_stream(inn, data)
359
+ end
360
+ end
361
+
362
+ raise_command_failed(status, cmd) unless status.success?
363
+
364
+ result[out].join.chomp
365
+
366
+ ensure
367
+ inn.close rescue nil
368
+ out.close rescue nil
369
+ err.close rescue nil
370
+ end
371
+
372
+
373
+ private
374
+
375
+ def raise_command_failed(status, cmd)
376
+ raise CmdError,
377
+ "Execution failed with status #{status.exitstatus}: #{[*cmd].join ' '}"
378
+ end
379
+
380
+
381
+ def password_required? stream_name, data
382
+ stream_name == :err && data =~ SUDO_PROMPT
383
+ end
384
+
385
+
386
+ def send_password_to_stream inn, data
387
+ prompt_for_password if data =~ SUDO_FAILED
388
+ inn.puts @password || prompt_for_password
389
+ end
390
+
391
+
392
+ def kill_process pid, kill_type="KILL"
393
+ begin
394
+ Process.kill kill_type, pid
395
+ Process.wait
396
+ rescue
397
+ end
398
+ end
399
+
400
+
401
+ def process_streams pid, *streams
402
+ result = Hash.new{|h,k| h[k] = []}
403
+ update_timeout
404
+
405
+ # Handle process termination ourselves
406
+ status = nil
407
+ Thread.start do
408
+ status = Process.waitpid2(pid).last
409
+ end
410
+
411
+ until streams.empty? do
412
+ # don't busy loop
413
+ selected, = select streams, nil, nil, 0.1
414
+
415
+ raise TimeoutError if timed_out?
416
+
417
+ next if selected.nil? or selected.empty?
418
+
419
+ selected.each do |stream|
420
+
421
+ update_timeout
422
+
423
+ if stream.eof? then
424
+ streams.delete stream if status # we've quit, so no more writing
425
+ next
426
+ end
427
+
428
+ data = stream.readpartial(1024)
429
+
430
+ yield(stream, data)
431
+
432
+ result[stream] << data
433
+ end
434
+ end
435
+
436
+ return result, status
437
+ end
438
+ end
439
+ end
440
+
@@ -0,0 +1,179 @@
1
+ require 'cgi'
2
+
3
+ module MockObject
4
+
5
+ ##
6
+ # Setup a method mock
7
+
8
+ def mock method, options={}, &block
9
+ mock_key = mock_key_for method, options
10
+ method_mocks[mock_key] = block_given? ? block : options[:return]
11
+ end
12
+
13
+
14
+ ##
15
+ # Get the value a mocked method was setup to return
16
+
17
+ def method_mock_return mock_key
18
+ return_val = method_mocks[mock_key] rescue method_mocks[[mock_key.first]]
19
+ if Proc === return_val
20
+ args = mock_key[1..-1]
21
+ return_val.call(*args)
22
+ else
23
+ return_val
24
+ end
25
+ end
26
+
27
+
28
+ ##
29
+ # Create a mock key based on :method, :args => [args_passed_to_method]
30
+
31
+ def mock_key_for method, options={}
32
+ mock_key = [method.to_s]
33
+ mock_key.concat [*options[:args]] if options.has_key?(:args)
34
+ mock_key
35
+ end
36
+
37
+
38
+ ##
39
+ # Check if a method was called. Supports options:
40
+ # :exactly:: num - exact number of times the method should have been called
41
+ # :count:: num - minimum number of times the method should have been called
42
+ # Defaults to :count => 1
43
+
44
+ def method_called? method, options={}
45
+ target_count = options[:count] || options[:exactly] || 1
46
+
47
+ count = method_call_count method, options
48
+
49
+ options[:exactly] ? count == target_count : count >= target_count
50
+ end
51
+
52
+
53
+ ##
54
+ # Count the number of times a method was called:
55
+ # obj.method_call_count :my_method, :args => [1,2,3]
56
+
57
+ def method_call_count method, options={}
58
+ count = 0
59
+
60
+ mock_def_arr = mock_key_for method, options
61
+
62
+ each_mock_key_matching(mock_def_arr) do |mock_key|
63
+ count = count + method_log[mock_key]
64
+ end
65
+
66
+ count
67
+ end
68
+
69
+
70
+ ##
71
+ # Do something with every instance of a mock key.
72
+ # Used to retrieve all lowest common denominators of method calls:
73
+ #
74
+ # each_mock_key_matching [:my_method] do |mock_key|
75
+ # puts mock_key.inspect
76
+ # end
77
+ #
78
+ # # Outputs #
79
+ # [:my_method, 1, 2, 3]
80
+ # [:my_method, 1, 2]
81
+ # [:my_method, 1]
82
+ # [:my_method]
83
+
84
+ def each_mock_key_matching mock_key
85
+ index = mock_key.length - 1
86
+
87
+ method_log.keys.each do |key|
88
+ yield(key) if block_given? && key[0..index] == mock_key
89
+ end
90
+ end
91
+
92
+
93
+ def method_mocks
94
+ @method_mocks ||= Hash.new do |h, k|
95
+ raise "Mock for #{k.inspect} does not exist."
96
+ end
97
+ end
98
+
99
+
100
+ def method_log
101
+ @method_log ||= Hash.new(0)
102
+ end
103
+
104
+
105
+ ##
106
+ # Hook into the object
107
+
108
+ def self.included base
109
+ hook_instance_methods base
110
+ end
111
+
112
+
113
+ def self.extended base
114
+ hook_instance_methods base, true
115
+ end
116
+
117
+
118
+ def self.hook_instance_methods base, instance=false
119
+ unhook_instance_methods base, instance
120
+
121
+ eval_each_method_of(base, instance) do |m|
122
+ m_def = m =~ /[^\]]=$/ ? "args" : "*args, &block"
123
+ new_m = escape_unholy_method_name "hooked_#{m}"
124
+ %{
125
+ alias #{new_m} #{m}
126
+ undef #{m}
127
+
128
+ def #{m}(#{m_def})
129
+ mock_key = mock_key_for '#{m}', :args => args
130
+
131
+ count = method_log[mock_key]
132
+ method_log[mock_key] = count.next
133
+
134
+ method_mock_return(mock_key) rescue self.send(:#{new_m}, #{m_def})
135
+ end
136
+ }
137
+ end
138
+ end
139
+
140
+
141
+ def self.unhook_instance_methods base, instance=false
142
+ eval_each_method_of(base, instance) do |m|
143
+ new_m = escape_unholy_method_name "hooked_#{m}"
144
+ #puts m + " -> " + new_m
145
+ %{
146
+ m = '#{new_m}'.to_sym
147
+ defined = method_defined?(m) rescue self.class.method_defined?(m)
148
+
149
+ if defined
150
+ undef #{m}
151
+ alias #{m} #{new_m}
152
+ end
153
+ }
154
+ end
155
+ end
156
+
157
+
158
+ def self.escape_unholy_method_name name
159
+ CGI.escape(name).gsub('%','').gsub('-','MNS')
160
+ end
161
+
162
+
163
+ def self.eval_each_method_of base, instance=false, &block
164
+ eval_method, affect_methods = if instance
165
+ [:instance_eval, base.methods]
166
+ else
167
+ [:class_eval, base.instance_methods]
168
+ end
169
+
170
+ banned_methods = self.instance_methods
171
+ banned_methods.concat Object.instance_methods
172
+
173
+ affect_methods.sort.each do |m|
174
+ next if banned_methods.include?(m)
175
+ #puts m
176
+ base.send eval_method, block.call(m)
177
+ end
178
+ end
179
+ end