add context manager to connect to a network share

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")
 

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)

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

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)