From b61773c7a21bfba8de69ae5ff995c6cbb0ffd306 Mon Sep 17 00:00:00 2001 From: Derek Hirotsu Date: Thu, 17 Jul 2025 16:31:00 -0700 Subject: [PATCH] add context manager to connect to a network share --- handbrake/handbrake.py | 18 ++++++++++++++++-- handbrake/handbrake_utils.py | 19 +++++++++++++++++++ handbrake/manifest.yaml | 6 ++++++ harness_utils/network_share.py | 23 +++++++++++++++++++++++ 4 files changed, 64 insertions(+), 2 deletions(-) create mode 100644 harness_utils/network_share.py diff --git a/handbrake/handbrake.py b/handbrake/handbrake.py index e1cd0ea..0f95f76 100644 --- a/handbrake/handbrake.py +++ b/handbrake/handbrake.py @@ -2,7 +2,15 @@ from argparse import ArgumentParser import os import re -from handbrake_utils import HANDBRAKE_EXECUTABLE, current_time_ms, handbrake_present, is_video_source_present, copy_video_source, copy_handbrake_from_network_drive +from handbrake_utils import ( + HANDBRAKE_EXECUTABLE, + current_time_ms, + handbrake_present, + is_video_source_present, + copy_video_source, + copy_handbrake_from_network_drive, + connect_and_copy_handbrake +) import logging import subprocess import sys @@ -94,6 +102,8 @@ def main(): parser = ArgumentParser() parser.add_argument("-e", "--encoder", dest="encoder", help="encoder", metavar="encoder", required=True) + parser.add_argument("--shareUser", dest="share_user", required=False, default=None) + parser.add_argument("--sharePass", dest="share_pass", required=False, default=None) args = parser.parse_args() if args.encoder not in ENCODER_TO_PRESET: @@ -105,7 +115,11 @@ def main(): preset = ENCODER_TO_PRESET[args.encoder] if handbrake_present() is False: logging.info("copying handbrake from network drive") - copy_handbrake_from_network_drive() + if args.share_user is not None and args.share_name is not None: + connect_and_copy_handbrake(args.share_user, args.share_pass) + else: + copy_handbrake_from_network_drive() + else: logging.info("detected handbrake") diff --git a/handbrake/handbrake_utils.py b/handbrake/handbrake_utils.py index 0169612..3a33e97 100644 --- a/handbrake/handbrake_utils.py +++ b/handbrake/handbrake_utils.py @@ -9,6 +9,9 @@ HANDBRAKE_EXECUTABLE = "HandBrakeCLI.exe" SOURCE_VIDEO_NAME = "big_buck_bunny_1080p24.y4m" SCRIPT_DIR = Path(os.path.dirname(os.path.realpath(__file__))) +from harness_utils.network_share import ( + network_share_auth +) def handbrake_present() -> bool: """Check if handbrake is present on the system""" @@ -23,6 +26,14 @@ def copy_handbrake_from_network_drive(): shutil.copyfile(copy_souce, destination) +def connect_and_copy_handbrake(username, password): + """Use context manager to connect to network drive and copy handbrake CLI""" + source = r"H:\01_Installers_Utilities\Handbrake\X86\HandBrakeCLI-1.9.1-win-x86_64\HandBrakeCLI.exe" + destination = SCRIPT_DIR / HANDBRAKE_EXECUTABLE + with network_share_auth(r"\\Labs\labs", username, password): + shutil.copyfile(source, destination) + + def is_video_source_present() -> bool: """check if big buck bunny video source is present""" return os.path.isfile(Path(SCRIPT_DIR / SOURCE_VIDEO_NAME)) @@ -36,6 +47,14 @@ def copy_video_source(): shutil.copyfile(source, destination) +def connect_and_copy_video(username, password): + """Use context manager to connect to network drive and copy video source""" + source = r"H:\03_ProcessingFiles\Handbrake Test\big_buck_bunny_1080p24.y4m" + destination = SCRIPT_DIR / SOURCE_VIDEO_NAME + with network_share_auth(r"\\Labs\labs", username, password): + shutil.copyfile(source, destination) + + def current_time_ms(): """Get current timestamp in milliseconds since epoch""" return int(time.time() * 1000) diff --git a/handbrake/manifest.yaml b/handbrake/manifest.yaml index 74089d2..4af1289 100644 --- a/handbrake/manifest.yaml +++ b/handbrake/manifest.yaml @@ -18,3 +18,9 @@ options: - h264_quicksync - av1_quicksync tooltip: Select which preset to use + - name: shareUser + type: input + tooltip: Username for connecting to network drive to retrive test files + - name: sharePass + type: input + tooltip: Password for connecting to network drive to retrive test files diff --git a/harness_utils/network_share.py b/harness_utils/network_share.py new file mode 100644 index 0000000..7650f17 --- /dev/null +++ b/harness_utils/network_share.py @@ -0,0 +1,23 @@ +""" +Provides a Contenxt Manager to connect to a network share +as a workaround to copying files directly from a pre-mounted network share. +Taken from this stackoverflow answer: https://stackoverflow.com/a/2626085 +""" +from contextlib import contextmanager +import os + +@contextmanager +def network_share_auth(share, username=None, password=None, drive_letter='H'): + """Context manager that mounts the given share using the given + username and password to the given drive letter when entering + the context and unmounts it when exiting.""" + cmd_parts = ["NET USE %s: %s" % (drive_letter, share)] + if password: + cmd_parts.append(password) + if username: + cmd_parts.append("/USER:%s" % username) + os.system(" ".join(cmd_parts)) + try: + yield + finally: + os.system("NET USE %s: /DELETE" % drive_letter)