def extract_pdf(self):
if self.file_location.startswith('\\'):
try:
with smbclient.open_file(r"{}".format(self.file_location),
mode='rb',
username=smb_username,
password=smb_password) as f:
pdf_extracted = pdfplumber.open(f)
print('file found')
except:
smbclient.reset_connection_cache()
with smbclient.open_file(r"{}".format(self.file_location),
mode='rb',
username=smb_username,
password=smb_password) as f:
pdf_extracted = pdfplumber.open(f)
print('file found')
finally:
smbclient.reset_connection_cache()
else:
file = document_location + self.file_location
print(file)
pdf_extracted = pdfplumber.open(file)
return pdf_extracted
def extract_cells(self, method=None): # extract cells on iterating sheets
if self.file_location.startswith('\\'):
try:
with smbclient.open_file(r"{}".format(self.file_location),
mode='rb',
username=smb_username,
password=smb_password) as f:
excel = pd.ExcelFile(f)
print('file found')
except:
smbclient.reset_connection_cache()
with smbclient.open_file(r"{}".format(self.file_location),
mode='rb',
username=smb_username,
password=smb_password) as f:
excel = pd.ExcelFile(f)
print('file found')
finally:
smbclient.reset_connection_cache()
else:
file = document_location + self.file_location
print(file)
excel = pd.ExcelFile(file)
# excel = pd.ExcelFile(self.file_location)
for sheet in excel.sheet_names:
self.data_frame = pd.read_excel(file,
header=None,
sheet_name=sheet)
rows, columns = self.data_frame.shape
for column in range(columns):
for row in range(rows):
yield column, row
def get_input(input_pdf, input_pdf_location):
if input_pdf.startswith('\\'):
print('connecting to SMB share')
try:
with smbclient.open_file(r"{}".format(input_pdf),
mode='rb',
username=smb_username,
password=smb_password) as f:
with open(input_pdf_location, 'wb') as pdf:
pdf.write(f.read())
print('file found')
except:
smbclient.reset_connection_cache()
with smbclient.open_file(r"{}".format(input_pdf),
mode='rb',
username=smb_username,
password=smb_password) as f:
with open(input_pdf_location, 'wb') as pdf:
pdf.write(f.read())
print('file found')
finally:
smbclient.reset_connection_cache()
return input_pdf_location
else:
return document_location + input_pdf
def upload(self,localFile,remotePath):
try:
smbclient.register_session(self.host, username=self.username, password=self.password, connection_timeout=30)
except Exception as e:
self.error = str(e)
return False
# single file
if not os.path.isdir(localFile):
try:
with open(localFile, mode="rb") as f:
with smbclient.open_file("\\{0}\{1}".format(self.host,remotePath), mode="wb") as remoteFile:
while True:
part = f.read(4096)
if not part:
break
remoteFile.write(part)
except Exception as e:
self.error = str(e)
smbclient.delete_session(self.host)
return False
smbclient.delete_session(self.host)
return True
# Directory
else:
try:
smbclient.mkdir("\\{0}\{1}".format(self.host,remotePath))
except OSError as e:
if e.errno != errno.EEXIST:
return False
for root, dirs, files in os.walk(localFile):
for dir in dirs:
fullPath = os.path.join(root,dir)
fullPath=fullPath.replace("/","\\")
try:
smbclient.mkdir("\\{0}\{1}\{2}".format(self.host,remotePath,fullPath[len(localFile)+1:]))
except OSError as e:
if e.errno != errno.EEXIST:
return False
for _file in files:
try:
fullPath = os.path.join(root,_file)
with open(fullPath, mode="rb")as f:
fullPath=fullPath.replace("/","\\")
with smbclient.open_file("\\{0}\{1}\{2}".format(self.host,remotePath,fullPath[len(localFile)+1:]), mode="wb") as remoteFile:
while True:
part = f.read(4096)
if not part:
break
remoteFile.write(part)
except Exception as e:
self.error = str(e)
smbclient.delete_session(self.host)
return False
smbclient.delete_session(self.host)
return True
return False
def test_samefile_with_different_files(smb_share):
file1 = "%s\\file1.txt" % smb_share
file2 = "%s\\file2.txt" % smb_share
with open_file(file1, mode='w') as fd:
fd.write(u"content")
with open_file(file2, mode='w') as fd:
fd.write(u"content")
assert samefile(file1, file2) is False
def get_file_smb(file_name):
data = None
try:
data = smbclient.open_file(r"{}".format(file_name),
mode='rb',
username=smb_username,
password=smb_password)
print('file found')
except:
smbclient.reset_connection_cache()
data = smbclient.open_file(r"{}".format(file_name), mode='rb')
finally:
return data
def fetch_xxxx(self,remotedri,filetpye,filestring = "."):
print('Start Download smbfile')
logging.info('Start Download smbfile')
Files = listdir(remotedri)
TotalFiles = 0
SuccessFiles = 0
for onefile in Files :
try :
filecompile = re.compile('{}'.format(filestring))
findlength = len(filecompile.findall(onefile))
if onefile.lower().endswith('.{}'.format(filetpye)) and findlength > 0:
Remotepath = os.path.join(remotedri,onefile)
Localpath = './temporaryfile/{}'.format(onefile)
TotalFiles+= 1
"""通过创建好的SMB共享盘客服端，读取远端文件内容写入本地存储。Remotepath要读取的远程文件路径，Localpath是要写入内容的本地文件路径"""
with open_file(Remotepath,mode = "rb") as fr:
file_bytes = fr.read()
with open(Localpath,"wb") as fw :
fw.write(file_bytes)
SuccessFiles+= 1
except Exception as inst :
print('Download file Fail')
print(inst)
print("smb download Total Number of Files: "+ str(TotalFiles) +"; Success Files: " +str(SuccessFiles))
logging.info("smb download Total Number of Files: "+ str(TotalFiles) +"; Success Files: " +str(SuccessFiles))
def openFile(self, path: str) -> SMBBackendFile:
obj = smbclient.open_file(
path, mode="rb", username=self._username, password=self._password
)
return SMBBackendFile(obj)
def executeCommand(self,command,args=[],elevate=False,runAs=None,timeout=300):
try:
if runAs:
if runAs == "user":
if args:
command = "{0}\nrm c:\\windows\\temp\\jimiRunAsUser.ps1".format(powershellRunAsLoggedInUser.replace("<CMD>",command))
else:
command = "{0}\nrm c:\\windows\\temp\\jimiRunAsUser.ps1".format(powershellRunAsLoggedInUser.replace("<CMD>",command))
with smbclient.open_file("\\{0}\{1}".format(self.host,"c$\\windows\\temp\\jimiRunAsUser.ps1"), mode="wb") as remoteFile:
for line in command.split("\n"):
remoteFile.write("{0}\n".format(line).encode())
clientShell = self.client.open_shell()
commandId = self.client.run_command(clientShell,"powershell.exe -noProfile -ExecutionPolicy Bypass -File c:\\windows\\temp\\jimiRunAsUser.ps1")
stdout, stderr, exitCode = self.client.get_command_output(clientShell, commandId)
self.client.cleanup_command(clientShell, commandId)
self.client.close_shell(clientShell)
response = stdout.decode().strip()
errors = stderr.decode().strip()
return (exitCode, response, errors)
else:
clientShell = self.client.open_shell()
commandId = self.client.run_command(clientShell,command, args)
stdout, stderr, exitCode = self.client.get_command_output(clientShell, commandId)
self.client.cleanup_command(clientShell, commandId)
self.client.close_shell(clientShell)
response = stdout.decode().strip()
errors = stderr.decode().strip()
return (exitCode, response, errors)
except Exception as e:
self.error = str(e)
return (None, None, None)
def _open(
self,
path,
mode="rb",
block_size=-1,
autocommit=True,
cache_options=None,
**kwargs,
):
"""
block_size: int or None
If 0, no buffering, 1, line buffering, >1, buffer that many bytes
Notes
-----
By specifying 'share_access' in 'kwargs' it is possible to override the
default shared access setting applied in the constructor of this object.
"""
bls = block_size if block_size is not None and block_size >= 0 else -1
wpath = _as_unc_path(self.host, path)
share_access = kwargs.pop("share_access", self.share_access)
if "w" in mode and autocommit is False:
temp = _as_temp_path(self.host, path, self.temppath)
return SMBFileOpener(wpath, temp, mode, block_size=bls, **kwargs)
return smbclient.open_file(
wpath, mode, buffering=bls, share_access=share_access, **kwargs
)
def download(self, remoteFile, localPath, createMissingFolders):
try:
smbclient.register_session(self.host,
username=self.username,
password=self.password,
connection_timeout=30)
except Exception as e:
self.error = str(e)
return False
try:
if createMissingFolders:
splitChar = "\\"
if splitChar not in localPath:
splitChar = "/"
if not os.path.isdir(localPath.rsplit(splitChar, 1)[0]):
os.makedirs(localPath.rsplit(splitChar, 1)[0])
with open(localPath, mode="wb") as f:
with smbclient.open_file("\\{0}\{1}".format(
self.host, remoteFile),
mode="rb") as remoteFile:
while True:
part = remoteFile.read(4096)
if not part:
break
f.write(part)
smbclient.delete_session(self.host)
return True
except Exception as e:
self.error = str(e)
smbclient.delete_session(self.host)
return False
def upload_for_smb(root, paths, target_path, checkpoint):
if len(paths) == 0:
return
for path in paths:
print('send path : ' + path)
try:
smbclient.mkdir(target_path + '/' + path,username=username, password=password, port=smb_port)
except:
pass
send_files = get_event_files(root, path)
for send_file in send_files:
try:
smbclient.stat(target_path + '/' + path + '/' + send_file.split('/')[-1],username=username, password=password, port=smb_port)
except:
print('sending ' + send_file)
dest = smbclient.open_file(target_path+'/'+path+'/'+send_file.split('/')[-1], 'wb',username=username, password=password, port=smb_port)
src = open(send_file, 'rb')
while True:
temp = src.read(1024 * 4)
if len(temp) == 0:
break
dest.write(temp)
src.close()
dest.close()
make_checkpoint(paths, checkpoint)
def upload_file(self, remotedri, filetpye, filestring="."):
"""
param:remotedri --> "\\10.41.22.77\Plant3\Special\M360\MZ3I00\中台数据库-Ason\\"
filetpye --> "xlsx"
filestring = "." --> filestring = "Scrap" 该文档具有的唯一特征。
fun:上传指定目录下，指定类型、指定特征的一个文档，到SMB指定目录下。
return: 返回上传的文档名称 :-->"NPI_PRINFO1.xls" onefile(包含filetype)
"""
print('Start To Upload file ！')
logging.info('Start To Upload file ！')
for _, _, files in os.walk('./smbtemporaryfile'):
for onefile in files:
try:
filecompile = re.compile('{}'.format(filestring))
findlength = len(filecompile.findall(onefile))
if onefile.lower().endswith(
'.{}'.format(filetpye)) and findlength > 0:
Remotepath = remotedri + onefile
Localpath = './smbtemporaryfile/{0}'.format(onefile)
with open(Localpath, mode="rb") as fr:
file_bytes = fr.read()
with open_file(Remotepath, mode="wb") as fw:
fw.write(file_bytes)
except Exception as inst:
print('Fail To Upload Smbfile !')
print(inst)
print('Success To Upload Smbfile !')
logging.info('Success To Upload Smbfile !')
return onefile
def open_file(
self,
path,
mode="r",
buffering=-1,
encoding=None,
errors=None,
newline=None,
share_access=None,
desired_access=None,
file_attributes=None,
file_type="file",
):
return smbclient.open_file(
self._join_path(path),
mode=mode,
buffering=buffering,
encoding=encoding,
errors=errors,
newline=newline,
share_access=share_access,
desired_access=desired_access,
file_attributes=file_attributes,
file_type=file_type,
**self._conn_kwargs,
)
def test_getsize(smb_share):
filename = "%s\\file.txt" % smb_share
with open_file(filename, mode='wb') as fd:
fd.write(b"\x00\x01\x02\x03")
assert getsize(filename) == 4
def smb_upload(client: SMBClient, args: dict):
hostname = args.get('hostname')
path = handle_path(args.get('file_path'))
path = os.path.join(hostname or client.hostname, path)
username = args.get('username')
password = args.get('password')
entryID = args.get('entryID')
content = args.get('content')
if not entryID and not content:
raise DemistoException(
"You must provide a content to upload using one of the following arguments: content, entryID."
)
client.create_session(hostname, username, password)
# For the content argument - the input is text.
writing_mode = 'w'
if entryID:
file = demisto.getFilePath(entryID)
filePath = file['path']
writing_mode = 'wb'
with open(filePath, mode='rb') as f:
content = f.read()
with open_file(fr'{path}', mode=writing_mode) as file_obj:
file_obj.write(content)
return f'File {get_file_name(path)} was uploaded successfully'
def get_emails_data():
"""
Read csv from shared drive with employee emails
:return: dictionary of emails
"""
employee_emails = {}
smbclient.ClientConfig(
username=os.getenv("SHAREDDRIVE_USERNAME"),
password=os.getenv("SHAREDDRIVE_PASSWORD"),
)
emails_csv_path = os.getenv("SHAREDDRIVE_FILEPATH")
# variables stored in json objects do not work well with \, replace / in stored path with \
emails_csv = emails_csv_path.replace("/", "\\")
with smbclient.open_file(emails_csv, mode="r") as emails:
reader = csv.DictReader(emails)
data = [row for row in reader]
for row in data:
# check if employeeID exists for record
if row.get("employeeID"):
# all contractors have ID 999, this will not collect their emails
employee_emails[row.get("employeeID")] = {
"email": row.get("EmailAddress"),
"name": row.get("Name"),
}
return employee_emails
def test_getctime(smb_share):
filename = "%s\\file.txt" % smb_share
with open_file(filename, mode='w') as fd:
fd.write(u"content")
expected = stat(filename).st_ctime
assert getctime(filename) == expected
def test_lexists(smb_share):
filename = "%s\\file.txt" % smb_share
assert exists(filename) is False
with open_file(filename, mode='w') as fd:
fd.write(u"content")
assert lexists(filename) is True
def test_islink_with_file_link(smb_share):
src_filename = "%s\\file.txt" % smb_share
dst_filename = "%s\\link.txt" % smb_share
with open_file(src_filename, mode='w') as fd:
fd.write(u"content")
symlink(src_filename, dst_filename)
assert islink(dst_filename) is True
def get_from_smb(self):
'''get file downloaded from smb share'''
if self.input_document.startswith('\\'):
with smbclient.open_file(r"{}".format(self.input_document), mode='rb', username=smb_username,
password=smb_password) as f:
with open(self.location_to_store, 'wb') as pdf:
pdf.write(f.read())
return self.location_to_store
else:
return document_location + self.input_document
def test_exists_working_symlink(smb_share):
src_filename = "%s\\file.txt" % smb_share
dst_filename = "%s\\link.txt" % smb_share
with open_file(src_filename, mode='w') as fd:
fd.write(u"content")
symlink(src_filename, dst_filename)
assert exists(dst_filename) is True
def test_samefile_missing_path2(smb_share):
file1 = "%s\\file1.txt" % smb_share
file2 = "%s\\file2.txt" % smb_share
with open_file(file2, mode='w') as fd:
fd.write(u"content")
expected = "[NtStatus 0xc0000034] No such file or directory: "
with pytest.raises(OSError, match=re.escape(expected)):
samefile(file1, file2)
def test_samefile_with_symlink(smb_share):
file1 = "%s\\file1.txt" % smb_share
file2 = "%s\\file2.txt" % smb_share
with open_file(file1, mode='w') as fd:
fd.write(u"content")
symlink(file1, file2)
assert samefile(file1, file2) is True
def smb_download(client: SMBClient, args: dict):
hostname = args.get('hostname')
path = handle_path(args.get('file_path'))
path = os.path.join(hostname or client.hostname, path)
username = args.get('username')
password = args.get('password')
client.create_session(hostname, username, password)
with open_file(fr'{path}', mode="rb") as fd:
file_contents = fd.read()
file_name = get_file_name(path)
return fileResult(file_name, file_contents)
def copy_test_two(source_file: str, target_file: str) -> None:
# Register credentials
smbclient.register_session(server, username=user, password=password)
# Absolut file path to the target file
destination = r"\\{server}\FAKEWINDOWS\repository\client-files\{file}".format(
server=server, file=target_file)
# Read source file
file_object = open(source_file, "r").read()
# Write target file
with smbclient.open_file(destination, mode="w") as remote_file:
remote_file.write(file_object)
def get_file(file):
if file.startswith('\\'):
print('connecting to SMB share')
try:
with smbclient.open_file(r"{}".format(file), mode='rb', username=smb_username,
password=smb_password) as f:
output_file.write(f.read())
output_file.seek(0)
print('file found')
except:
raise Exception('No access for the given path')
finally:
smbclient.reset_connection_cache()
return 'SMB'
else:
return 'LOCAL'
def _open_file_smbprotocol(url, mode="rb", **kwargs):
domain, host, port, user, passwd, server_path = _parse_smb_url(url)
import smbclient
try:
# register the server with explicit credentials
if user:
session = smbclient.register_session(
host, username=user, password=passwd, port=port
)
# Read an existing file as bytes
mode2 = mode[:1] + r"b"
filehandle = smbclient.open_file(server_path, mode=mode2, **kwargs)
return filehandle
except Exception as ex:
raise ConnectionError("SMB error: %s" % ex).with_traceback(sys.exc_info()[2])
def get_smb_or_local(input_pdf, input_pdf_location):
if input_pdf.startswith('\\'):
print('connecting to SMB share')
try:
with smbclient.open_file(r"{}".format(input_pdf),
mode='rb',
username=smb_username,
password=smb_password) as f:
with open(input_pdf_location, 'wb') as pdf:
pdf.write(f.read())
print('file found')
except:
pass
# raise Exception('File is being used by another process / smb not accessible')
finally:
smbclient.reset_connection_cache()
return input_pdf_location
else:
return document_location + input_pdf
def download(self,remoteFile,localPath,createMissingFolders):
try:
if createMissingFolders:
splitChar = "\\"
if splitChar not in localPath:
splitChar = "/"
if not os.path.isdir(localPath.rsplit(splitChar,1)[0]):
os.makedirs(localPath.rsplit(splitChar,1)[0])
with open(localPath, mode="wb") as f:
with smbclient.open_file("\\{0}\{1}".format(self.host,remoteFile), mode="rb") as remoteFile:
while True:
part = remoteFile.read(4096)
if not part:
break
f.write(part)
return True
except Exception as e:
self.error = str(e)
return False