Skip to content

Commit

Permalink
Configurable post-restore script and time synchronization
Browse files Browse the repository at this point in the history
  • Loading branch information
psrok1 committed Sep 23, 2024
1 parent f69281a commit 9417f1b
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 17 deletions.
52 changes: 35 additions & 17 deletions drakrun/drakrun/analyzer.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import argparse
import base64
import contextlib
import dataclasses
import datetime
import itertools
import json
import logging
Expand Down Expand Up @@ -173,25 +175,41 @@ def drop_sample_to_vm(
raise e


def prepare_post_restore_envs(options: AnalysisOptions):
current_date = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
return {
"$DRAKVUF_NET_ENABLE": "$true" if options.net_enable else "$false",
"$DRAKVUF_DATE": f'"{current_date}"',
}


def prepare_post_restore_command(options: AnalysisOptions):
post_restore_script_path = pathlib.Path(ETC_DIR) / "scripts" / "vm-post-restore.ps1"
post_restore_script = post_restore_script_path.read_text()
post_restore_envs = prepare_post_restore_envs(options)
for env_key, env_value in post_restore_envs.items():
post_restore_script = post_restore_script.replace(env_key, env_value)
log.debug("Prepared post-restore script: " + post_restore_script)
encoded_ps1 = base64.b64encode(post_restore_script.encode("utf-16le")).decode()
return "powershell.exe -EncodedCommand " + encoded_ps1


def run_post_restore_vm_commands(vm: DrakvufVM, options: AnalysisOptions):
log.info("Running post-restore VM commands...")
if options.net_enable:
max_attempts = 3
for i in range(max_attempts):
try:
log.info(f"Trying to setup network (attempt {i + 1}/{max_attempts})")
vm.injector.create_process(
"cmd /C ipconfig /release >nul", wait=True, timeout=120
)
vm.injector.create_process(
"cmd /C ipconfig /renew >nul", wait=True, timeout=120
)
break
except Exception:
log.exception("Analysis attempt failed. Retrying...")
else:
log.warning(f"Giving up after {max_attempts} failures...")
raise RuntimeError("Failed to setup VM network after 3 attempts")
max_attempts = 3
for i in range(max_attempts):
try:
log.info(
f"Trying to run post-restore commands (attempt {i + 1}/{max_attempts})"
)
post_restore_command = prepare_post_restore_command(options)
vm.injector.create_process(post_restore_command, wait=True, timeout=120)
break
except Exception:
log.exception("Post-restore command execution failed. Retrying...")
else:
log.warning(f"Giving up after {max_attempts} failures...")
raise RuntimeError("Failed to run post-restore commands after 3 attempts")


def get_profile_args_list() -> List[str]:
Expand Down
6 changes: 6 additions & 0 deletions drakrun/drakrun/data/vm-post-restore.ps1
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
if ( $DRAKVUF_NET_ENABLE )
{
ipconfig /release
ipconfig /renew
}
Set-Date -Date $DRAKVUF_DATE
3 changes: 3 additions & 0 deletions drakrun/drakrun/draksetup/init.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,9 @@ def is_component_to_init(component_name):
"cfg.template", target_dir=ETC_SCRIPTS_DIR
)
set_template_vnc_password(template_path)
create_configuration_file(
"vm-post-restore.ps1", target_dir=ETC_SCRIPTS_DIR
)

if is_component_to_init("system"):
create_configuration_file(
Expand Down

0 comments on commit 9417f1b

Please sign in to comment.