secator 0.5.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

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.

Potentially problematic release.


This version of secator might be problematic. Click here for more details.

Files changed (84) hide show
  1. secator/celery.py +160 -185
  2. secator/celery_utils.py +268 -0
  3. secator/cli.py +327 -106
  4. secator/config.py +27 -11
  5. secator/configs/workflows/host_recon.yaml +5 -3
  6. secator/configs/workflows/port_scan.yaml +7 -3
  7. secator/configs/workflows/url_bypass.yaml +10 -0
  8. secator/configs/workflows/url_vuln.yaml +1 -1
  9. secator/decorators.py +169 -92
  10. secator/definitions.py +10 -3
  11. secator/exporters/__init__.py +7 -5
  12. secator/exporters/console.py +10 -0
  13. secator/exporters/csv.py +27 -19
  14. secator/exporters/gdrive.py +16 -11
  15. secator/exporters/json.py +3 -1
  16. secator/exporters/table.py +30 -2
  17. secator/exporters/txt.py +20 -16
  18. secator/hooks/gcs.py +53 -0
  19. secator/hooks/mongodb.py +54 -28
  20. secator/output_types/__init__.py +29 -11
  21. secator/output_types/_base.py +11 -1
  22. secator/output_types/error.py +36 -0
  23. secator/output_types/exploit.py +1 -1
  24. secator/output_types/info.py +24 -0
  25. secator/output_types/ip.py +7 -0
  26. secator/output_types/port.py +8 -1
  27. secator/output_types/progress.py +6 -1
  28. secator/output_types/record.py +3 -1
  29. secator/output_types/stat.py +33 -0
  30. secator/output_types/tag.py +6 -4
  31. secator/output_types/url.py +6 -3
  32. secator/output_types/vulnerability.py +3 -2
  33. secator/output_types/warning.py +24 -0
  34. secator/report.py +55 -23
  35. secator/rich.py +44 -39
  36. secator/runners/_base.py +622 -635
  37. secator/runners/_helpers.py +5 -91
  38. secator/runners/celery.py +18 -0
  39. secator/runners/command.py +364 -211
  40. secator/runners/scan.py +8 -24
  41. secator/runners/task.py +21 -55
  42. secator/runners/workflow.py +41 -40
  43. secator/scans/__init__.py +28 -0
  44. secator/serializers/dataclass.py +6 -0
  45. secator/serializers/json.py +10 -5
  46. secator/serializers/regex.py +12 -4
  47. secator/tasks/_categories.py +6 -3
  48. secator/tasks/bbot.py +293 -0
  49. secator/tasks/bup.py +98 -0
  50. secator/tasks/cariddi.py +38 -49
  51. secator/tasks/dalfox.py +3 -0
  52. secator/tasks/dirsearch.py +12 -23
  53. secator/tasks/dnsx.py +49 -30
  54. secator/tasks/dnsxbrute.py +2 -0
  55. secator/tasks/feroxbuster.py +8 -17
  56. secator/tasks/ffuf.py +3 -2
  57. secator/tasks/fping.py +3 -3
  58. secator/tasks/gau.py +5 -0
  59. secator/tasks/gf.py +2 -2
  60. secator/tasks/gospider.py +4 -0
  61. secator/tasks/grype.py +9 -9
  62. secator/tasks/h8mail.py +31 -41
  63. secator/tasks/httpx.py +58 -21
  64. secator/tasks/katana.py +18 -22
  65. secator/tasks/maigret.py +26 -24
  66. secator/tasks/mapcidr.py +2 -3
  67. secator/tasks/msfconsole.py +4 -16
  68. secator/tasks/naabu.py +3 -1
  69. secator/tasks/nmap.py +50 -35
  70. secator/tasks/nuclei.py +9 -2
  71. secator/tasks/searchsploit.py +17 -9
  72. secator/tasks/subfinder.py +5 -1
  73. secator/tasks/wpscan.py +79 -93
  74. secator/template.py +61 -45
  75. secator/thread.py +24 -0
  76. secator/utils.py +330 -80
  77. secator/utils_test.py +48 -23
  78. secator/workflows/__init__.py +28 -0
  79. {secator-0.5.2.dist-info → secator-0.7.0.dist-info}/METADATA +12 -6
  80. secator-0.7.0.dist-info/RECORD +115 -0
  81. {secator-0.5.2.dist-info → secator-0.7.0.dist-info}/WHEEL +1 -1
  82. secator-0.5.2.dist-info/RECORD +0 -101
  83. {secator-0.5.2.dist-info → secator-0.7.0.dist-info}/entry_points.txt +0 -0
  84. {secator-0.5.2.dist-info → secator-0.7.0.dist-info}/licenses/LICENSE +0 -0
@@ -1,20 +1,24 @@
1
+ import copy
1
2
  import getpass
2
3
  import logging
3
4
  import os
4
5
  import re
5
6
  import shlex
7
+ import signal
6
8
  import subprocess
7
9
  import sys
10
+ import uuid
8
11
 
9
- from time import sleep
12
+ from time import time
10
13
 
14
+ import psutil
11
15
  from fp.fp import FreeProxy
12
16
 
13
- from secator.template import TemplateLoader
14
17
  from secator.definitions import OPT_NOT_SUPPORTED, OPT_PIPE_INPUT
15
18
  from secator.config import CONFIG
19
+ from secator.output_types import Info, Error, Target, Stat
16
20
  from secator.runners import Runner
17
- from secator.serializers import JSONSerializer
21
+ from secator.template import TemplateLoader
18
22
  from secator.utils import debug
19
23
 
20
24
 
@@ -53,18 +57,18 @@ class Command(Runner):
53
57
  # Output encoding
54
58
  encoding = 'utf-8'
55
59
 
56
- # Environment variables
57
- env = {}
58
-
59
60
  # Flag to take the input
60
61
  input_flag = None
61
62
 
62
63
  # Input path (if a file is constructed)
63
64
  input_path = None
64
65
 
65
- # Input chunk size (default None)
66
+ # Input chunk size
66
67
  input_chunk_size = CONFIG.runners.input_chunk_size
67
68
 
69
+ # Input required
70
+ input_required = True
71
+
68
72
  # Flag to take a file as input
69
73
  file_flag = None
70
74
 
@@ -80,7 +84,14 @@ class Command(Runner):
80
84
 
81
85
  # Serializer
82
86
  item_loader = None
83
- item_loaders = [JSONSerializer(),]
87
+ item_loaders = []
88
+
89
+ # Hooks
90
+ hooks = [
91
+ 'on_cmd',
92
+ 'on_cmd_done',
93
+ 'on_line'
94
+ ]
84
95
 
85
96
  # Ignore return code
86
97
  ignore_return_code = False
@@ -88,9 +99,6 @@ class Command(Runner):
88
99
  # Return code
89
100
  return_code = -1
90
101
 
91
- # Error
92
- error = ''
93
-
94
102
  # Output
95
103
  output = ''
96
104
 
@@ -102,7 +110,8 @@ class Command(Runner):
102
110
  # Profile
103
111
  profile = 'cpu'
104
112
 
105
- def __init__(self, input=None, **run_opts):
113
+ def __init__(self, inputs=[], **run_opts):
114
+
106
115
  # Build runnerconfig on-the-fly
107
116
  config = TemplateLoader(input={
108
117
  'name': self.__class__.__name__,
@@ -110,26 +119,45 @@ class Command(Runner):
110
119
  'description': run_opts.get('description', None)
111
120
  })
112
121
 
113
- # Run parent init
122
+ # Extract run opts
114
123
  hooks = run_opts.pop('hooks', {})
124
+ caller = run_opts.get('caller', None)
115
125
  results = run_opts.pop('results', [])
116
126
  context = run_opts.pop('context', {})
127
+ self.skip_if_no_inputs = run_opts.pop('skip_if_no_inputs', False)
128
+
129
+ # Prepare validators
130
+ input_validators = []
131
+ if not self.skip_if_no_inputs:
132
+ input_validators.append(self._validate_input_nonempty)
133
+ if not caller:
134
+ input_validators.append(self._validate_chunked_input)
135
+ validators = {'validate_input': input_validators}
136
+
137
+ # Call super().__init__
117
138
  super().__init__(
118
139
  config=config,
119
- targets=input,
140
+ inputs=inputs,
120
141
  results=results,
121
142
  run_opts=run_opts,
122
143
  hooks=hooks,
144
+ validators=validators,
123
145
  context=context)
124
146
 
147
+ # Inputs path
148
+ self.inputs_path = None
149
+
125
150
  # Current working directory for cmd
126
151
  self.cwd = self.run_opts.get('cwd', None)
127
152
 
128
- # No capturing of stdout / stderr.
129
- self.no_capture = self.run_opts.get('no_capture', False)
153
+ # Print cmd
154
+ self.print_cmd = self.run_opts.get('print_cmd', False)
155
+
156
+ # Stat update
157
+ self.last_updated_stat = None
130
158
 
131
- # No processing of output lines.
132
- self.no_process = self.run_opts.get('no_process', False)
159
+ # Process
160
+ self.process = None
133
161
 
134
162
  # Proxy config (global)
135
163
  self.proxy = self.run_opts.pop('proxy', False)
@@ -141,6 +169,9 @@ class Command(Runner):
141
169
  # Build command
142
170
  self._build_cmd()
143
171
 
172
+ # Run on_cmd hook
173
+ self.run_hooks('on_cmd')
174
+
144
175
  # Build item loaders
145
176
  instance_func = getattr(self, 'item_loader', None)
146
177
  item_loaders = self.item_loaders.copy()
@@ -148,40 +179,6 @@ class Command(Runner):
148
179
  item_loaders.append(instance_func)
149
180
  self.item_loaders = item_loaders
150
181
 
151
- # Print built cmd
152
- if self.print_cmd and not self.has_children:
153
- if self.sync and self.description:
154
- self._print(f'\n:wrench: {self.description} ...', color='bold gold3', rich=True)
155
- self._print(self.cmd.replace('[', '\\['), color='bold cyan', rich=True)
156
-
157
- # Print built input
158
- if self.print_input_file and self.input_path:
159
- input_str = '\n '.join(self.input).strip()
160
- debug(f'[dim magenta]File input:[/]\n [italic medium_turquoise]{input_str}[/]')
161
-
162
- # Print run options
163
- if self.print_run_opts:
164
- input_str = '\n '.join([
165
- f'[dim blue]{k}[/] -> [dim green]{v}[/]' for k, v in self.run_opts.items() if v is not None]).strip()
166
- debug(f'[dim magenta]Run opts:[/]\n {input_str}')
167
-
168
- # Print format options
169
- if self.print_fmt_opts:
170
- input_str = '\n '.join([
171
- f'[dim blue]{k}[/] -> [dim green]{v}[/]' for k, v in self.opts_to_print.items() if v is not None]).strip()
172
- debug(f'[dim magenta]Print opts:[/]\n {input_str}')
173
-
174
- # Print hooks
175
- if self.print_hooks:
176
- input_str = ''
177
- for hook_name, hook_funcs in self.hooks.items():
178
- hook_funcs_str = ', '.join([f'[dim green]{h.__module__}.{h.__qualname__}[/]' for h in hook_funcs])
179
- if hook_funcs:
180
- input_str += f'[dim blue]{hook_name}[/] -> {hook_funcs_str}\n '
181
- input_str = input_str.strip()
182
- if input_str:
183
- debug(f'[dim magenta]Hooks:[/]\n {input_str}')
184
-
185
182
  def toDict(self):
186
183
  res = super().toDict()
187
184
  res.update({
@@ -196,6 +193,7 @@ class Command(Runner):
196
193
  # TODO: Move this to TaskBase
197
194
  from secator.celery import run_command
198
195
  results = kwargs.get('results', [])
196
+ kwargs['sync'] = False
199
197
  name = cls.__name__
200
198
  return run_command.apply_async(args=[results, name] + list(args), kwargs={'opts': kwargs}, queue=cls.profile)
201
199
 
@@ -206,7 +204,7 @@ class Command(Runner):
206
204
  return run_command.s(cls.__name__, *args, opts=kwargs).set(queue=cls.profile)
207
205
 
208
206
  @classmethod
209
- def si(cls, results, *args, **kwargs):
207
+ def si(cls, *args, results=[], **kwargs):
210
208
  # TODO: Move this to TaskBase
211
209
  from secator.celery import run_command
212
210
  return run_command.si(results, cls.__name__, *args, opts=kwargs).set(queue=cls.profile)
@@ -226,12 +224,14 @@ class Command(Runner):
226
224
  d[k] = v.__name__
227
225
  return d
228
226
 
229
- opts = {k: convert(v) for k, v in cls.opts.items()}
227
+ cls_opts = copy.deepcopy(cls.opts)
228
+ opts = {k: convert(v) for k, v in cls_opts.items()}
230
229
  for k, v in opts.items():
231
230
  v['meta'] = cls.__name__
232
231
  v['supported'] = True
233
232
 
234
- meta_opts = {k: convert(v) for k, v in cls.meta_opts.items() if cls.opt_key_map.get(k) is not OPT_NOT_SUPPORTED}
233
+ cls_meta_opts = copy.deepcopy(cls.meta_opts)
234
+ meta_opts = {k: convert(v) for k, v in cls_meta_opts.items() if cls.opt_key_map.get(k) is not OPT_NOT_SUPPORTED}
235
235
  for k, v in meta_opts.items():
236
236
  v['meta'] = 'meta'
237
237
  if cls.opt_key_map.get(k) is OPT_NOT_SUPPORTED:
@@ -247,7 +247,7 @@ class Command(Runner):
247
247
  #---------------#
248
248
 
249
249
  @classmethod
250
- def execute(cls, cmd, name=None, cls_attributes={}, **kwargs):
250
+ def execute(cls, cmd, name=None, cls_attributes={}, run=True, **kwargs):
251
251
  """Execute an ad-hoc command.
252
252
 
253
253
  Can be used without defining an inherited class to run a command, while still enjoying all the good stuff in
@@ -264,15 +264,13 @@ class Command(Runner):
264
264
  secator.runners.Command: instance of the Command.
265
265
  """
266
266
  name = name or cmd.split(' ')[0]
267
- kwargs['no_process'] = kwargs.get('no_process', True)
268
267
  kwargs['print_cmd'] = not kwargs.get('quiet', False)
269
- kwargs['print_item'] = not kwargs.get('quiet', False)
270
- kwargs['print_line'] = not kwargs.get('quiet', False)
271
- delay_run = kwargs.pop('delay_run', False)
272
- cmd_instance = type(name, (Command,), {'cmd': cmd})(**kwargs)
268
+ kwargs['print_line'] = True
269
+ kwargs['no_process'] = kwargs.get('no_process', True)
270
+ cmd_instance = type(name, (Command,), {'cmd': cmd, 'input_required': False})(**kwargs)
273
271
  for k, v in cls_attributes.items():
274
272
  setattr(cmd_instance, k, v)
275
- if not delay_run:
273
+ if run:
276
274
  cmd_instance.run()
277
275
  return cmd_instance
278
276
 
@@ -328,124 +326,228 @@ class Command(Runner):
328
326
 
329
327
  Yields:
330
328
  str: Command stdout / stderr.
331
- dict: Parsed JSONLine object.
329
+ dict: Serialized object.
332
330
  """
333
- # Set status to 'RUNNING'
334
- self.status = 'RUNNING'
331
+ try:
335
332
 
336
- # Callback before running command
337
- self.run_hooks('on_start')
333
+ # Abort if it has children tasks
334
+ if self.has_children:
335
+ return
338
336
 
339
- # Check for sudo requirements and prepare the password if needed
340
- sudo_password = self._prompt_sudo(self.cmd)
337
+ # Print task description
338
+ self.print_description()
341
339
 
342
- # Prepare cmds
343
- command = self.cmd if self.shell else shlex.split(self.cmd)
340
+ # Abort if no inputs
341
+ if len(self.inputs) == 0 and self.skip_if_no_inputs:
342
+ yield Info(message=f'{self.unique_name} skipped (no inputs)', _source=self.unique_name, _uuid=str(uuid.uuid4()))
343
+ return
344
344
 
345
- # Output and results
346
- self.return_code = 0
347
- self.killed = False
345
+ # Yield targets
346
+ for input in self.inputs:
347
+ yield Target(name=input, _source=self.unique_name, _uuid=str(uuid.uuid4()))
348
+
349
+ # Check for sudo requirements and prepare the password if needed
350
+ sudo_password, error = self._prompt_sudo(self.cmd)
351
+ if error:
352
+ yield Error(
353
+ message=error,
354
+ _source=self.unique_name,
355
+ _uuid=str(uuid.uuid4())
356
+ )
357
+ return
348
358
 
349
- # Run the command using subprocess
350
- try:
359
+ # Prepare cmds
360
+ command = self.cmd if self.shell else shlex.split(self.cmd)
361
+
362
+ # Output and results
363
+ self.return_code = 0
364
+ self.killed = False
365
+
366
+ # Run the command using subprocess
351
367
  env = os.environ
352
- env.update(self.env)
353
368
  self.process = subprocess.Popen(
354
369
  command,
355
370
  stdin=subprocess.PIPE if sudo_password else None,
356
- stdout=sys.stdout if self.no_capture else subprocess.PIPE,
357
- stderr=sys.stderr if self.no_capture else subprocess.STDOUT,
371
+ stdout=subprocess.PIPE,
372
+ stderr=subprocess.STDOUT,
358
373
  universal_newlines=True,
359
374
  shell=self.shell,
360
375
  env=env,
361
376
  cwd=self.cwd)
377
+ self.print_command()
362
378
 
363
379
  # If sudo password is provided, send it to stdin
364
380
  if sudo_password:
365
381
  self.process.stdin.write(f"{sudo_password}\n")
366
382
  self.process.stdin.flush()
367
383
 
368
- except FileNotFoundError as e:
369
- if self.config.name in str(e):
370
- error = 'Executable not found.'
371
- if self.install_cmd:
372
- error += f' Install it with `secator install tools {self.config.name}`.'
373
- else:
374
- error = str(e)
375
- celery_id = self.context.get('celery_id', '')
376
- if celery_id:
377
- error += f' [{celery_id}]'
378
- self.errors.append(error)
379
- self.return_code = 1
380
- return
381
-
382
- try:
383
- # No capture mode, wait for command to finish and return
384
- if self.no_capture:
385
- self._wait_for_end()
386
- return
387
-
388
384
  # Process the output in real-time
389
385
  for line in iter(lambda: self.process.stdout.readline(), b''):
390
- sleep(0) # for async to give up control
386
+ # sleep(0) # for async to give up control
391
387
  if not line:
392
388
  break
389
+ yield from self.process_line(line)
393
390
 
394
- # Strip line endings
395
- line = line.rstrip()
396
- if self.no_process:
397
- yield line
398
- continue
391
+ # Run hooks after cmd has completed successfully
392
+ result = self.run_hooks('on_cmd_done')
393
+ if result:
394
+ yield from result
395
+
396
+ except FileNotFoundError as e:
397
+ yield from self.handle_file_not_found(e)
398
+
399
+ except BaseException as e:
400
+ self.debug(f'{self.unique_name}: {type(e).__name__}.', sub='error')
401
+ self.stop_process()
402
+ yield Error.from_exception(e, _source=self.unique_name, _uuid=str(uuid.uuid4()))
403
+
404
+ finally:
405
+ yield from self._wait_for_end()
406
+
407
+ def process_line(self, line):
408
+ """Process a single line of output emitted on stdout / stderr and yield results."""
409
+
410
+ # Strip line endings
411
+ line = line.rstrip()
412
+
413
+ # Some commands output ANSI text, so we need to remove those ANSI chars
414
+ if self.encoding == 'ansi':
415
+ ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
416
+ line = ansi_escape.sub('', line)
417
+ line = line.replace('\\x0d\\x0a', '\n')
418
+
419
+ # Run on_line hooks
420
+ line = self.run_hooks('on_line', line)
421
+ if line is None:
422
+ return
399
423
 
400
- # Some commands output ANSI text, so we need to remove those ANSI chars
401
- if self.encoding == 'ansi':
402
- # ansi_regex = r'\x1b\[([0-9,A-Z]{1,2}(;[0-9]{1,2})?(;[0-9]{3})?)?[K]?'
403
- # line = re.sub(ansi_regex, '', line.strip())
404
- ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
405
- line = ansi_escape.sub('', line)
406
- line = line.replace('\\x0d\\x0a', '\n')
424
+ # Run item_loader to try parsing as dict
425
+ item_count = 0
426
+ for item in self.run_item_loaders(line):
427
+ yield item
428
+ item_count += 1
407
429
 
408
- # Run on_line hooks
409
- line = self.run_hooks('on_line', line)
430
+ # Yield line if no items were yielded
431
+ if item_count == 0:
432
+ yield line
410
433
 
411
- # Run item_loader to try parsing as dict
412
- items = None
413
- if self.output_json:
414
- items = self.run_item_loaders(line)
434
+ # Skip rest of iteration (no process mode)
435
+ if self.no_process:
436
+ return
415
437
 
416
- # Yield line if no items parsed
417
- if not items:
418
- yield line
438
+ # Yield command stats (CPU, memory, conns ...)
439
+ # TODO: enable stats support with timer
440
+ if self.last_updated_stat and (time() - self.last_updated_stat) < CONFIG.runners.stat_update_frequency:
441
+ return
419
442
 
420
- # Turn results into list if not already a list
421
- elif not isinstance(items, list):
422
- items = [items]
443
+ yield from self.stats()
444
+ self.last_updated_stat = time()
445
+
446
+ def print_description(self):
447
+ """Print description"""
448
+ if self.sync and not self.has_children:
449
+ if self.caller and self.description:
450
+ self._print(f'\n[bold gold3]:wrench: {self.description} [dim cyan]({self.config.name})[/][/] ...', rich=True)
451
+ elif self.print_cmd:
452
+ self._print('')
453
+
454
+ def print_command(self):
455
+ """Print command."""
456
+ if self.print_cmd:
457
+ cmd_str = self.cmd.replace('[', '\\[')
458
+ if self.sync and self.chunk and self.chunk_count:
459
+ cmd_str += f' [dim gray11]({self.chunk}/{self.chunk_count})[/]'
460
+ self._print(cmd_str, color='bold cyan', rich=True)
461
+ self.debug('Command', obj={'cmd': self.cmd}, sub='init')
462
+
463
+ def handle_file_not_found(self, exc):
464
+ """Handle case where binary is not found.
423
465
 
424
- # Yield items
425
- if items:
426
- yield from items
466
+ Args:
467
+ exc (FileNotFoundError): the exception.
427
468
 
428
- except KeyboardInterrupt:
429
- self.process.kill()
430
- self.killed = True
469
+ Yields:
470
+ secator.output_types.Error: the error.
471
+ """
472
+ self.return_code = 127
473
+ if self.config.name in str(exc):
474
+ message = 'Executable not found.'
475
+ if self.install_cmd:
476
+ message += f' Install it with `secator install tools {self.config.name}`.'
477
+ error = Error(message=message)
478
+ else:
479
+ error = Error.from_exception(exc)
480
+ error._source = self.unique_name
481
+ error._uuid = str(uuid.uuid4())
482
+ yield error
483
+
484
+ def stop_process(self):
485
+ """Sends SIGINT to running process, if any."""
486
+ if not self.process:
487
+ return
488
+ self.debug(f'Sending SIGINT to process {self.process.pid}.', sub='error')
489
+ self.process.send_signal(signal.SIGINT)
490
+
491
+ def stats(self):
492
+ """Gather stats about the current running process, if any."""
493
+ if not self.process or not self.process.pid:
494
+ return
495
+ proc = psutil.Process(self.process.pid)
496
+ stats = Command.get_process_info(proc, children=True)
497
+ for info in stats:
498
+ name = info['name']
499
+ pid = info['pid']
500
+ cpu_percent = info['cpu_percent']
501
+ mem_percent = info['memory_percent']
502
+ net_conns = info.get('net_connections') or []
503
+ extra_data = {k: v for k, v in info.items() if k not in ['cpu_percent', 'memory_percent', 'net_connections']}
504
+ yield Stat(
505
+ name=name,
506
+ pid=pid,
507
+ cpu=cpu_percent,
508
+ memory=mem_percent,
509
+ net_conns=len(net_conns),
510
+ extra_data=extra_data
511
+ )
431
512
 
432
- # Retrieve the return code and output
433
- self._wait_for_end()
513
+ @staticmethod
514
+ def get_process_info(process, children=False):
515
+ """Get process information from psutil.
516
+
517
+ Args:
518
+ process (subprocess.Process): Process.
519
+ children (bool): Whether to gather stats about children processes too.
520
+ """
521
+ try:
522
+ data = {
523
+ k: v._asdict() if hasattr(v, '_asdict') else v
524
+ for k, v in process.as_dict().items()
525
+ if k not in ['memory_maps', 'open_files', 'environ']
526
+ }
527
+ yield data
528
+ except (psutil.Error, FileNotFoundError):
529
+ return
530
+ if children:
531
+ for subproc in process.children(recursive=True):
532
+ yield from Command.get_process_info(subproc, children=False)
434
533
 
435
534
  def run_item_loaders(self, line):
436
- """Run item loaders on a string."""
437
- items = []
535
+ """Run item loaders against an output line.
536
+
537
+ Args:
538
+ line (str): Output line.
539
+ """
540
+ if self.no_process:
541
+ return
438
542
  for item_loader in self.item_loaders:
439
- result = None
440
543
  if (callable(item_loader)):
441
- result = item_loader(self, line)
544
+ yield from item_loader(self, line)
442
545
  elif item_loader:
443
- result = item_loader.run(line)
444
- if isinstance(result, dict):
445
- result = [result]
446
- if result:
447
- items.extend(result)
448
- return items
546
+ name = item_loader.__class__.__name__.replace('Serializer', '').lower()
547
+ default_callback = lambda self, x: [(yield x)] # noqa: E731
548
+ callback = getattr(self, f'on_{name}_loaded', None) or default_callback
549
+ for item in item_loader.run(line):
550
+ yield from callback(self, item)
449
551
 
450
552
  def _prompt_sudo(self, command):
451
553
  """
@@ -455,22 +557,25 @@ class Command(Runner):
455
557
  command (str): The initial command to be executed.
456
558
 
457
559
  Returns:
458
- str or None: The sudo password if required; otherwise, None.
560
+ tuple: (sudo password, error).
459
561
  """
460
562
  sudo_password = None
461
563
 
462
564
  # Check if sudo is required by the command
463
565
  if not re.search(r'\bsudo\b', command):
464
- return None
566
+ return None, []
465
567
 
466
568
  # Check if sudo can be executed without a password
467
- if subprocess.run(['sudo', '-n', 'true'], capture_output=True).returncode == 0:
468
- return None
569
+ try:
570
+ if subprocess.run(['sudo', '-n', 'true'], capture_output=False).returncode == 0:
571
+ return None, None
572
+ except ValueError:
573
+ self._print('[bold orange3]Could not run sudo check test.[/][bold green]Passing.[/]')
469
574
 
470
575
  # Check if we have a tty
471
576
  if not os.isatty(sys.stdin.fileno()):
472
- self._print("No TTY detected. Sudo password prompt requires a TTY to proceed.", color='bold red')
473
- sys.exit(1)
577
+ error = "No TTY detected. Sudo password prompt requires a TTY to proceed."
578
+ return -1, error
474
579
 
475
580
  # If not, prompt the user for a password
476
581
  self._print('[bold red]Please enter sudo password to continue.[/]')
@@ -484,31 +589,43 @@ class Command(Runner):
484
589
  capture_output=True
485
590
  )
486
591
  if result.returncode == 0:
487
- return sudo_password # Password is correct
592
+ return sudo_password, None # Password is correct
488
593
  self._print("Sorry, try again.")
489
- self._print("Sudo password verification failed after 3 attempts.")
490
- return None
594
+ error = "Sudo password verification failed after 3 attempts."
595
+ return -1, error
491
596
 
492
597
  def _wait_for_end(self):
493
598
  """Wait for process to finish and process output and return code."""
599
+ if not self.process:
600
+ return
601
+ for line in self.process.stdout.readlines():
602
+ yield from self.process_line(line)
494
603
  self.process.wait()
495
604
  self.return_code = self.process.returncode
605
+ self.process.stdout.close()
606
+ self.return_code = 0 if self.ignore_return_code else self.return_code
607
+ self.output = self.output.strip()
608
+ self.killed = self.return_code == -2 or self.killed
609
+ self.debug(f'Command {self.cmd} finished with return code {self.return_code}', sub='command')
496
610
 
497
- if self.no_capture:
498
- self.output = ''
499
- else:
500
- self.output = self.output.strip()
501
- self.process.stdout.close()
502
-
503
- if self.ignore_return_code:
504
- self.return_code = 0
505
-
506
- if self.return_code == -2 or self.killed:
611
+ if self.killed:
507
612
  error = 'Process was killed manually (CTRL+C / CTRL+X)'
508
- self.errors.append(error)
613
+ yield Error(
614
+ message=error,
615
+ _source=self.unique_name,
616
+ _uuid=str(uuid.uuid4())
617
+ )
618
+
509
619
  elif self.return_code != 0:
510
620
  error = f'Command failed with return code {self.return_code}.'
511
- self.errors.append(error)
621
+ last_lines = self.output.split('\n')
622
+ last_lines = last_lines[max(0, len(last_lines) - 2):]
623
+ yield Error(
624
+ message=error,
625
+ traceback='\n'.join(last_lines),
626
+ _source=self.unique_name,
627
+ _uuid=str(uuid.uuid4())
628
+ )
512
629
 
513
630
  @staticmethod
514
631
  def _process_opts(
@@ -518,15 +635,19 @@ class Command(Runner):
518
635
  opt_value_map={},
519
636
  opt_prefix='-',
520
637
  command_name=None):
521
- """Process a dict of options using a config, option key map / value map
522
- and option character like '-' or '--'.
638
+ """Process a dict of options using a config, option key map / value map and option character like '-' or '--'.
523
639
 
524
640
  Args:
525
641
  opts (dict): Command options as input on the CLI.
526
642
  opts_conf (dict): Options config (Click options definition).
643
+ opt_key_map (dict[str, str | Callable]): A dict to map option key with their actual values.
644
+ opt_value_map (dict, str | Callable): A dict to map option values with their actual values.
645
+ opt_prefix (str, default: '-'): Option prefix.
646
+ command_name (str | None, default: None): Command name.
527
647
  """
528
648
  opts_str = ''
529
649
  for opt_name, opt_conf in opts_conf.items():
650
+ debug('before get_opt_value', obj={'name': opt_name, 'conf': opt_conf}, obj_after=False, sub='command.options', verbose=True) # noqa: E501
530
651
 
531
652
  # Get opt value
532
653
  default_val = opt_conf.get('default')
@@ -537,25 +658,30 @@ class Command(Runner):
537
658
  opt_prefix=command_name,
538
659
  default=default_val)
539
660
 
661
+ debug('after get_opt_value', obj={'name': opt_name, 'value': opt_val, 'conf': opt_conf}, obj_after=False, sub='command.options', verbose=True) # noqa: E501
662
+
540
663
  # Skip option if value is falsy
541
664
  if opt_val in [None, False, []]:
542
- # logger.debug(f'Option {opt_name} was passed but is falsy. Skipping.')
665
+ debug('skipped (falsy)', obj={'name': opt_name, 'value': opt_val}, obj_after=False, sub='command.options', verbose=True) # noqa: E501
543
666
  continue
544
667
 
545
668
  # Convert opt value to expected command opt value
546
669
  mapped_opt_val = opt_value_map.get(opt_name)
547
- if callable(mapped_opt_val):
548
- opt_val = mapped_opt_val(opt_val)
549
- elif mapped_opt_val:
550
- opt_val = mapped_opt_val
670
+ if mapped_opt_val:
671
+ if callable(mapped_opt_val):
672
+ opt_val = mapped_opt_val(opt_val)
673
+ else:
674
+ opt_val = mapped_opt_val
551
675
 
552
676
  # Convert opt name to expected command opt name
553
677
  mapped_opt_name = opt_key_map.get(opt_name)
554
- if mapped_opt_name == OPT_NOT_SUPPORTED:
555
- # logger.debug(f'Option {opt_name} was passed but is unsupported. Skipping.')
556
- continue
557
- elif mapped_opt_name is not None:
558
- opt_name = mapped_opt_name
678
+ if mapped_opt_name is not None:
679
+ if mapped_opt_name == OPT_NOT_SUPPORTED:
680
+ debug('skipped (unsupported)', obj={'name': opt_name, 'value': opt_val}, sub='command.options', verbose=True) # noqa: E501
681
+ continue
682
+ else:
683
+ opt_name = mapped_opt_name
684
+ debug('mapped key / value', obj={'name': opt_name, 'value': opt_val}, obj_after=False, sub='command.options', verbose=True) # noqa: E501
559
685
 
560
686
  # Avoid shell injections and detect opt prefix
561
687
  opt_name = str(opt_name).split(' ')[0] # avoid cmd injection
@@ -570,31 +696,66 @@ class Command(Runner):
570
696
  # Append opt name + opt value to option string.
571
697
  # Note: does not append opt value if value is True (flag)
572
698
  opts_str += f' {opt_name}'
699
+ shlex_quote = opt_conf.get('shlex', True)
573
700
  if opt_val is not True:
574
- opt_val = shlex.quote(str(opt_val))
701
+ if shlex_quote:
702
+ opt_val = shlex.quote(str(opt_val))
575
703
  opts_str += f' {opt_val}'
704
+ debug('final', obj={'name': opt_name, 'value': opt_val}, sub='command.options', obj_after=False, verbose=True)
576
705
 
577
706
  return opts_str.strip()
578
707
 
708
+ @staticmethod
709
+ def _validate_chunked_input(self, inputs):
710
+ """Command does not suport multiple inputs in non-worker mode. Consider using .delay() instead."""
711
+ if len(inputs) > 1 and self.sync and self.file_flag is None:
712
+ return False
713
+ return True
714
+
715
+ @staticmethod
716
+ def _validate_input_nonempty(self, inputs):
717
+ """Input is empty."""
718
+ if not self.input_required:
719
+ return True
720
+ if not inputs or len(inputs) == 0:
721
+ return False
722
+ return True
723
+
724
+ # @staticmethod
725
+ # def _validate_input_types_valid(self, input):
726
+ # pass
727
+
728
+ @staticmethod
729
+ def _get_opt_default(opt_name, opts_conf):
730
+ for k, v in opts_conf.items():
731
+ if k == opt_name:
732
+ return v.get('default', None)
733
+ return None
734
+
579
735
  @staticmethod
580
736
  def _get_opt_value(opts, opt_name, opts_conf={}, opt_prefix='', default=None):
581
- aliases = [
582
- opts.get(f'{opt_prefix}_{opt_name}'),
583
- opts.get(f'{opt_prefix}.{opt_name}'),
584
- opts.get(opt_name),
737
+ default = default or Command._get_opt_default(opt_name, opts_conf)
738
+ opt_names = [
739
+ f'{opt_prefix}.{opt_name}',
740
+ f'{opt_prefix}_{opt_name}',
741
+ opt_name,
585
742
  ]
586
- alias = [conf.get('short') for _, conf in opts_conf.items() if conf.get('short') in opts]
743
+ opt_values = [opts.get(o) for o in opt_names]
744
+ alias = [conf.get('short') for _, conf in opts_conf.items() if conf.get('short') in opts and _ == opt_name]
587
745
  if alias:
588
- aliases.append(opts.get(alias[0]))
589
- if OPT_NOT_SUPPORTED in aliases:
746
+ opt_values.append(opts.get(alias[0]))
747
+ if OPT_NOT_SUPPORTED in opt_values:
748
+ debug('skipped (unsupported)', obj={'name': opt_name}, obj_after=False, sub='command.options', verbose=True)
590
749
  return None
591
- return next((v for v in aliases if v is not None), default)
750
+ value = next((v for v in opt_values if v is not None), default)
751
+ debug('got opt value', obj={'name': opt_name, 'value': value, 'aliases': opt_names, 'values': opt_values}, obj_after=False, sub='command.options', verbose=True) # noqa: E501
752
+ return value
592
753
 
593
754
  def _build_cmd(self):
594
755
  """Build command string."""
595
756
 
596
757
  # Add JSON flag to cmd
597
- if self.output_json and self.json_flag:
758
+ if self.json_flag:
598
759
  self.cmd += f' {self.json_flag}'
599
760
 
600
761
  # Add options to cmd
@@ -624,47 +785,39 @@ class Command(Runner):
624
785
  string or a list to the cmd.
625
786
  """
626
787
  cmd = self.cmd
627
- input = self.input
788
+ inputs = self.inputs
628
789
 
629
- # If input is None, return the previous command
630
- if not input:
790
+ # If inputs is empty, return the previous command
791
+ if not inputs:
631
792
  return
632
793
 
633
- # If input is a list but has one element, use the standard string input
634
- if isinstance(input, list) and len(input) == 1:
635
- input = input[0]
794
+ # If inputs has a single element but the tool does not support an input flag, use echo-piped_input input.
795
+ # If the tool's input flag is set to None, assume it is a positional argument at the end of the command.
796
+ # Otherwise use the input flag to pass the input.
797
+ if len(inputs) == 1:
798
+ input = shlex.quote(inputs[0])
799
+ if self.input_flag == OPT_PIPE_INPUT:
800
+ cmd = f'echo {input} | {cmd}'
801
+ elif not self.input_flag:
802
+ cmd += f' {input}'
803
+ else:
804
+ cmd += f' {self.input_flag} {input}'
636
805
 
637
- # If input is a list and the tool has input_flag set to OPT_PIPE_INPUT, use cat-piped input.
806
+ # If inputs has multiple elements and the tool has input_flag set to OPT_PIPE_INPUT, use cat-piped_input input.
638
807
  # Otherwise pass the file path to the tool.
639
- if isinstance(input, list):
808
+ else:
640
809
  fpath = f'{self.reports_folder}/.inputs/{self.unique_name}.txt'
641
810
 
642
811
  # Write the input to a file
643
812
  with open(fpath, 'w') as f:
644
- f.write('\n'.join(input))
813
+ f.write('\n'.join(inputs))
645
814
 
646
815
  if self.file_flag == OPT_PIPE_INPUT:
647
816
  cmd = f'cat {fpath} | {cmd}'
648
817
  elif self.file_flag:
649
818
  cmd += f' {self.file_flag} {fpath}'
650
- else:
651
- self._print(f'{self.__class__.__name__} does not support multiple inputs.', color='bold red')
652
- self.input_valid = False
653
-
654
- self.input_path = fpath
655
819
 
656
- # If input is a string but the tool does not support an input flag, use echo-piped input.
657
- # If the tool's input flag is set to None, assume it is a positional argument at the end of the command.
658
- # Otherwise use the input flag to pass the input.
659
- else:
660
- input = shlex.quote(input)
661
- if self.input_flag == OPT_PIPE_INPUT:
662
- cmd = f'echo {input} | {cmd}'
663
- elif not self.input_flag:
664
- cmd += f' {input}'
665
- else:
666
- cmd += f' {self.input_flag} {input}'
820
+ self.inputs_path = fpath
667
821
 
668
822
  self.cmd = cmd
669
823
  self.shell = ' | ' in self.cmd
670
- self.input = input