Your Name
241 lines
9.3 KiB
Python
241 lines
9.3 KiB
Python
# Copyright (c) 2024 Broadcom. All Rights Reserved.
|
|
# The term "Broadcom" refers to Broadcom Inc.
|
|
# and/or its subsidiaries.
|
|
|
|
import base64
|
|
import json
|
|
import glob
|
|
import os
|
|
import subprocess
|
|
from collections.abc import Mapping
|
|
from datetime import datetime
|
|
|
|
from lib.environment import Environment
|
|
from lib.exceptions import ReplayEntryNotFound, CommandExecutionError
|
|
|
|
replay_context = None
|
|
|
|
|
|
class ReplayEntry(object):
|
|
"""
|
|
Class for representing execution result entry
|
|
"""
|
|
|
|
def __init__(self, category, command_args, input_string, tag):
|
|
"""
|
|
ReplayEntry constructor
|
|
|
|
:param category: category of result (e.g. 'command', 'ldap')
|
|
:param command_args: command and the parameters
|
|
:param input_string: Input text to be supplied as stdin
|
|
"""
|
|
self.category = category
|
|
self.command_args = command_args
|
|
self.input_string = input_string if input_string is not None else ''
|
|
self.used_count = 0
|
|
self.key = self.generate_key(category, command_args, input_string, tag)
|
|
self.results = []
|
|
|
|
def add_result(self, timestamp, return_code, stdout, stderr):
|
|
"""
|
|
Add command result entry
|
|
|
|
:param timestamp: the timestamp of command execution
|
|
:param return_code: command's return code
|
|
:param stdout: The command output from stdout
|
|
:param stderr: The command output from stderr
|
|
"""
|
|
# always add result in sorted
|
|
added = False
|
|
value = (timestamp, return_code, stdout, stderr)
|
|
for index, (ts, _, _, _) in enumerate(self.results):
|
|
if ts > timestamp:
|
|
self.results.insert(index, value)
|
|
added = True
|
|
break
|
|
if not added:
|
|
self.results.append(value)
|
|
|
|
@staticmethod
|
|
def generate_key(category, command_args, input_string, tag):
|
|
"""
|
|
Generate key based on category, command_args, and input_string.
|
|
The key is not guaranteed to be unique but should be good enough for
|
|
debugging and testing purpose. This key will be used to retrieve
|
|
the previous command result for replay
|
|
|
|
:param category: the category, 'command', 'ldap'
|
|
:param command_args: command arguments
|
|
:param input_string: input string
|
|
:param tag: additional tag to differentiate the same commands in a different condition
|
|
:return: generated key
|
|
"""
|
|
for idx, arg in enumerate(command_args):
|
|
if not isinstance(arg, str):
|
|
command_args[idx] = str(arg)
|
|
if input_string == subprocess.DEVNULL:
|
|
input_string = '__/dev/null__'
|
|
return "{}||{}||{}||{}".format(category, ">>".join(command_args), input_string, tag)
|
|
|
|
|
|
class ReplayContext(object):
|
|
"""
|
|
The class holding context object for capturing or replaying command result
|
|
"""
|
|
|
|
categories = ['command', 'ldap', 'other']
|
|
|
|
def __init__(self, replay_dir, is_replaying=True, is_capturing=False):
|
|
"""
|
|
ReplayContext constructor
|
|
|
|
:param replay_dir: base directory for storing or retrieving command result
|
|
:param is_capturing: True when capturing, otherwise it's in replay mode
|
|
"""
|
|
self.replay_dir = replay_dir
|
|
self.replay_entries = None
|
|
self.is_replaying = is_replaying
|
|
self.is_capturing = is_capturing
|
|
self.tag = None
|
|
|
|
def get_replay_entry(self, key) -> ReplayEntry:
|
|
"""
|
|
Get replay entry based on the key
|
|
|
|
:param key: key to be used to retrieve the command result
|
|
"""
|
|
if self.replay_entries is None:
|
|
self.load_all_replays()
|
|
return self.replay_entries.get(key)
|
|
|
|
def store_result(self, category, command_args, command_input, return_code, stdout, stderr):
|
|
"""
|
|
Store the command result
|
|
|
|
The file will be stored in <replay-dir>/<hostname>/replay_<category>_<1st command_args>_<timestamp>.json
|
|
"""
|
|
command = command_args[0].split('/')[-1]
|
|
timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S_%f')
|
|
file_name = "{}/replay_{}_{}_{}.json".format(self.replay_dir, category,
|
|
command, timestamp)
|
|
with open(file_name, 'w') as file:
|
|
content = {
|
|
"command_args": ReplayContext.encode_bytes(command_args),
|
|
"input": command_input if command_input != subprocess.DEVNULL else '__/dev/null__',
|
|
"tag": self.tag,
|
|
"timestamp": timestamp,
|
|
"return_code": return_code,
|
|
"stdout": self.encode_bytes(stdout),
|
|
"stderr": self.encode_bytes(stderr)
|
|
}
|
|
try:
|
|
file.write(json.dumps(content, indent=4))
|
|
except TypeError as te:
|
|
raise CommandExecutionError("Unserializable object: {}".format(content))
|
|
|
|
if self.is_replaying:
|
|
self.add_replay_entry(category, command_args, command_input, self.tag, timestamp,
|
|
return_code, stdout, stderr)
|
|
|
|
def set_tag(self, tag):
|
|
self.tag = tag
|
|
|
|
def load_all_replays(self):
|
|
"""
|
|
Load all replay entries
|
|
"""
|
|
self.replay_entries = {}
|
|
for category in ReplayContext.categories:
|
|
for file in sorted(glob.glob("{}/replay_{}_*.json".format(self.replay_dir, category))):
|
|
with open(file, 'r') as f:
|
|
replay = json.load(f)
|
|
command_args = ReplayContext.decode_bytes(replay['command_args'])
|
|
command_input = replay['input']
|
|
stdout = ReplayContext.decode_bytes(replay['stdout'])
|
|
stderr = ReplayContext.decode_bytes(replay['stderr'])
|
|
self.add_replay_entry(category, command_args, command_input, replay['tag'], replay['timestamp'],
|
|
replay['return_code'], stdout, stderr)
|
|
|
|
def add_replay_entry(self, category, command_args, command_input, tag, timestamp,
|
|
return_code, stdout, stderr):
|
|
"""
|
|
Add replay entry
|
|
"""
|
|
key = ReplayEntry.generate_key(category, command_args, command_input, tag)
|
|
entry = self.get_replay_entry(key)
|
|
if entry is None:
|
|
entry = ReplayEntry(category, command_args, command_input, tag)
|
|
self.replay_entries[key] = entry
|
|
entry.add_result(timestamp, return_code, stdout, stderr)
|
|
|
|
def reset_context(self):
|
|
for entry in self.replay_entries:
|
|
entry.used_count = 0
|
|
|
|
def get_execution_result(self, category, command_args, command_input) -> (int, str, str):
|
|
"""
|
|
Get the previous execution result based on {category}, {command_args}, and {command_input}
|
|
"""
|
|
key = ReplayEntry.generate_key(category, command_args, command_input, self.tag)
|
|
entry = self.get_replay_entry(key)
|
|
if entry is None:
|
|
if self.is_capturing:
|
|
return None
|
|
else:
|
|
raise ReplayEntryNotFound("Replay entry not found for key {}".format(key))
|
|
result = entry.results[entry.used_count]
|
|
entry.used_count = (entry.used_count + 1) % len(entry.results)
|
|
return result[1:]
|
|
|
|
@staticmethod
|
|
def get_replay_context(reload=False):
|
|
"""
|
|
Static method to get ReplayContext instance from other modules
|
|
|
|
:param reload: if True, it will dispose the previous object and create a new instance
|
|
"""
|
|
global replay_context
|
|
env = Environment.get_environment()
|
|
is_remote = env.get_value('VCERT_REMOTE_EXEC') is True
|
|
is_replay = env.get_value('VCERT_REMOTE_EXEC_REPLAY') is True
|
|
is_capture = env.get_value('VCERT_REMOTE_EXEC_CAPTURE') is True
|
|
if reload:
|
|
replay_context = None
|
|
if is_remote and (is_replay or is_capture) and replay_context is None:
|
|
remote_hostname = env.get_value('VCERT_REMOTE_HOSTNAME')
|
|
replay_dir = "{}/{}".format(env.get_value('VCERT_REMOTE_EXEC_REPLAY_DIR'),
|
|
remote_hostname)
|
|
if not os.path.exists(replay_dir):
|
|
os.makedirs(replay_dir)
|
|
replay_context = ReplayContext(replay_dir, is_replay, is_capture)
|
|
return replay_context
|
|
|
|
@staticmethod
|
|
def encode_bytes(obj):
|
|
if isinstance(obj, bytes):
|
|
return "__base64__({})".format(base64.b64encode(obj).decode('utf-8'))
|
|
if isinstance(obj, list):
|
|
for idx, value in enumerate(obj):
|
|
obj[idx] = ReplayContext.encode_bytes(value)
|
|
elif isinstance(obj, dict):
|
|
for key in obj.keys():
|
|
obj[key] = ReplayContext.encode_bytes(obj[key])
|
|
elif isinstance(obj, Mapping):
|
|
new_obj = dict()
|
|
for key in obj.keys():
|
|
new_obj[key] = ReplayContext.encode_bytes(obj[key])
|
|
obj = new_obj
|
|
return obj
|
|
|
|
@staticmethod
|
|
def decode_bytes(obj):
|
|
if isinstance(obj, str) and obj.startswith('__base64__(') and obj.endswith(')'):
|
|
return base64.b64decode(obj[11:-1])
|
|
if isinstance(obj, list):
|
|
for idx, value in enumerate(obj):
|
|
obj[idx] = ReplayContext.decode_bytes(value)
|
|
elif isinstance(obj, dict):
|
|
for key in obj.keys():
|
|
obj[key] = ReplayContext.decode_bytes(obj[key])
|
|
return obj
|