This commit is contained in:
Ida Wang 2022-09-20 13:51:40 -04:00
Родитель 43baf65f84
Коммит d92ab5c729
4 изменённых файлов: 67 добавлений и 6 удалений

Просмотреть файл

@ -7,6 +7,7 @@ from __future__ import unicode_literals
import logging import logging
from datetime import datetime from datetime import datetime
import shutil
import threading import threading
import os import os
import sys import sys
@ -14,7 +15,6 @@ import glob
import pkgutil import pkgutil
import inspect import inspect
import importlib import importlib
import tempfile
import re import re
from Queue import Queue from Queue import Queue
@ -57,11 +57,15 @@ class AzureBatchAssets(object):
self._assets = None self._assets = None
self._tab_index = index self._tab_index = index
self._upload_threads = None self._upload_threads = None
self._temp_dir = utils.create_temp_dir()
self.batch = None self.batch = None
self.modules = self._collect_modules() self.modules = self._collect_modules()
self.ui = AssetsUI(self, frame) self.ui = AssetsUI(self, frame)
self.frame = frame self.frame = frame
def __del__(self):
shutil.rmtree(self._temp_dir)
def _callback_refresh(self, *args): def _callback_refresh(self, *args):
"""Called by Maya when a new scene file is loaded, so we reset """Called by Maya when a new scene file is loaded, so we reset
@ -160,7 +164,7 @@ class AzureBatchAssets(object):
:param str os_flavor: The chosen operating system of the render nodes, used :param str os_flavor: The chosen operating system of the render nodes, used
to determine the formatting of the path remapping. to determine the formatting of the path remapping.
""" """
proj_file = os.path.join(tempfile.gettempdir(), "workspace.mel") proj_file = os.path.join(self._temp_dir, "workspace.mel")
with open(proj_file, 'w') as handle: with open(proj_file, 'w') as handle:
for rule in maya.workspace(fileRuleList=True): for rule in maya.workspace(fileRuleList=True):
project_dir = maya.workspace(fileRuleEntry=rule) project_dir = maya.workspace(fileRuleEntry=rule)
@ -185,7 +189,7 @@ class AzureBatchAssets(object):
:param str os_flavor: The chosen operating system of the render nodes, used :param str os_flavor: The chosen operating system of the render nodes, used
to determine the formatting of the path remapping. to determine the formatting of the path remapping.
""" """
map_file = os.path.join(tempfile.gettempdir(), "asset_map.mel") map_file = os.path.join(self._temp_dir, "asset_map.mel")
pathmap = dict(self._assets.pathmaps) pathmap = dict(self._assets.pathmaps)
for asset in self._assets.refs: for asset in self._assets.refs:
pathmap.update(asset.pathmap) pathmap.update(asset.pathmap)

Просмотреть файл

@ -8,6 +8,7 @@ import os
import logging import logging
import platform import platform
import subprocess import subprocess
import tempfile
from azurebatchmayaapi import MayaAPI as maya from azurebatchmayaapi import MayaAPI as maya
from exception import CancellationException, FileUploadException from exception import CancellationException, FileUploadException
@ -145,6 +146,13 @@ def build_template_filename(render_engine, maya_version, operating_system, conta
return template_file return template_file
def create_temp_dir():
"""Create a temporary private local directory.
Should be deleted manually.
"""
return tempfile.mkdtemp()
class OperatingSystem(Enum): class OperatingSystem(Enum):
windows = 'Windows' windows = 'Windows'
linux = 'Linux' linux = 'Linux'

Просмотреть файл

@ -19,6 +19,7 @@ from azurebatchmayaapi import MayaAPI as maya
from ui_jobhistory import JobHistoryUI from ui_jobhistory import JobHistoryUI
import azure.batch as batch import azure.batch as batch
import azurebatchutils as utils
class AzureBatchJobHistory(object): class AzureBatchJobHistory(object):
@ -86,9 +87,9 @@ class AzureBatchJobHistory(object):
maya.refresh() maya.refresh()
return return
try: try:
temp_dir = os.path.join(tempfile.gettempdir(), job.id) temp_dir = os.path.join(utils.create_temp_dir(), job.id)
if not os.path.isdir(temp_dir): # if not os.path.isdir(temp_dir):
os.mkdir(temp_dir) # os.mkdir(temp_dir)
except Exception as exp: except Exception as exp:
self._log.warning(exp) self._log.warning(exp)
self.selected_job.set_thumbnail(thumb, self._get_image_height(thumb)) self.selected_job.set_thumbnail(thumb, self._get_image_height(thumb))

Просмотреть файл

@ -9,6 +9,17 @@ import logging
import json import json
from Queue import Queue from Queue import Queue
# win32-specific imports
try:
import pywintypes
import win32con
import win32net
import win32netcon
import win32security
import winerror
except ImportError:
pass
if sys.version_info >= (3, 3): if sys.version_info >= (3, 3):
import unittest2 as unittest import unittest2 as unittest
from unittest.mock import MagicMock from unittest.mock import MagicMock
@ -529,6 +540,43 @@ class TestAzureBatchAssets(unittest.TestCase):
self.mock_self._assets = mock.create_autospec(Assets) self.mock_self._assets = mock.create_autospec(Assets)
AzureBatchAssets.upload(self.mock_self) AzureBatchAssets.upload(self.mock_self)
@mock.patch.object(AzureBatchAssets, "_collect_modules")
@mock.patch("assets.callback")
@mock.patch("assets.AssetsUI")
def test_temp_dir_access(self, mock_ui, mock_call, mock_collect):
if sys.platform != "win32":
return
aba = AzureBatchAssets(0, "frame", "call")
# create temporary user
username = "tempuser"
password = "Password123"
user_info = {
"name": username,
"password": password,
"priv": win32netcon.USER_PRIV_USER,
"flags": 0
}
try:
win32net.NetUserAdd(None, 1, user_info)
except pywintypes.error as e:
if e.winerror == winerror.ERROR_ACCESS_DENIED:
raise Exception("test must be run as admin to create temp user")
else:
raise e
# make temp user admin
sid, domain, at = win32security.LookupAccountName(None, username)
win32net.NetLocalGroupAddMembers(None, "Administrators", 0, [{"sid": sid}])
# impersonate temporary user
handle = win32security.LogonUser(username, None, password, win32con.LOGON32_LOGON_INTERACTIVE, win32con.LOGON32_PROVIDER_DEFAULT)
win32security.ImpersonateLoggedOnUser(handle)
# try to access the temp directory
accessed = os.access(aba._temp_dir, os.W_OK)
# revert impersonation and delete temp user
win32security.RevertToSelf()
handle.close()
win32net.NetUserDel(None, username)
self.assertFalse(accessed, "_temp_dir should be inaccessible for other users")
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()