import re import shutil from os.path import isfile import demistomock as demisto import pyzipper from CommonServerPython import * from CommonServerUserPython import * ESCAPE_CHARACTERS = r'[/\<>"|?*]' def test_compression_succeeded(zip_name: str, password: str = None): with pyzipper.AESZipFile(zip_name) as zf: # testing for file integrity if password: zf.setpassword(bytes(password, "utf-8")) ret = zf.testzip() if ret is not None: demisto.info("zf.testzip() failed") raise DemistoException("There was a problem with zipping the file: " + ret + " is corrupted") def compress_multiple(file_names: List[str], zip_name: str, password: str = None): """ Compress multiple files into a zip file. :param file_names: list of file names to compress :param zip_name: name of the zip file to create :param password: password to use for encryption """ compression = pyzipper.ZIP_DEFLATED encryption = pyzipper.WZ_AES if password else None demisto.debug(f"zipping {file_names=}") with pyzipper.AESZipFile(zip_name, mode="w", compression=compression, encryption=encryption) as zf: zf.pwd = bytes(password, "utf-8") if password else None for file_name in file_names: zf.write(file_name) test_compression_succeeded(zip_name, password) zf.close() def escape_illegal_characters_in_file_name(file_name: str) -> str: if file_name: file_name = re.sub(ESCAPE_CHARACTERS, "-", file_name) file_name = re.sub(r"-+", "-", file_name) # prevent more than one consecutive dash in the file name return file_name def main(): try: args = demisto.args() zipName = None password = None fileEntryID = args.get("entryID") if "zipName" in args: zipName = escape_illegal_characters_in_file_name(demisto.args().get("zipName")) + ".zip" if "password" in args: password = demisto.args().get("password") if not fileEntryID: raise DemistoException("You must set an entryID when using the zip script") entry_ids = argToList(fileEntryID) file_names = [] for entry_id in entry_ids: res = demisto.executeCommand("getFilePath", {"id": entry_id}) if is_error(res): raise DemistoException( "Failed to get the file path for entry: " + entry_id + " the error message was " + get_error(res) ) filePath = res[0]["Contents"]["path"] fileCurrentName = escape_illegal_characters_in_file_name(res[0]["Contents"]["name"]) if not isfile(filePath): # in case that the user will send a directory raise DemistoException(entry_id + " is not a file. Please recheck your input.") # Handling duplicate names. if fileCurrentName in file_names: name, ext = os.path.splitext(fileCurrentName) i = 0 while fileCurrentName in file_names: i += 1 fileCurrentName = f"{name} {i}{ext}" # copying the file to current location shutil.copy(filePath, fileCurrentName) file_names.append(fileCurrentName) if not zipName: # Preserving old behaviour. If only one file provided - will use its name .zip # Else will use a uuid. if len(file_names) == 1: fileCurrentName = file_names[0] else: fileCurrentName = demisto.uniqueFile() zipName = fileCurrentName + ".zip" # zipping the file compress_multiple(file_names, zipName, password) with open(zipName, "rb") as f: file_data = f.read() demisto.results(fileResult(zipName, file_data)) human_readable = tableToMarkdown("Zipped Files", [{"original name": file_names, "zipped file": zipName}]) context: Dict[str, Any] = {"ZippedFiles": zipName, "ZipFile.ZippedFile": zipName} for entry_id in entry_ids: context[f"File(val.EntryID == {entry_id}).zipped"] = True raw_response = {"ZippedFiles": zipName} return_outputs(human_readable, context, raw_response) except Exception as exc: return_error(exc) if __name__ in ("__main__", "__builtin__", "builtins"): main()