import subprocess import demistomock as demisto from CommonServerPython import * TESSERACT_EXE = "tesseract" CORRUPTED_ERRORS = ("pix not read", "truncated file") CORRUPTED_ERR = CORRUPTED_ERRORS[0] # kept for backward compatibility CORRUPTED_MSG = "WARNING: failed to extract text - image is corrupted" EMPTY_FILE_MSG = "WARNING: failed to extract text - file is empty" def list_languages() -> list[str]: lang_out = subprocess.check_output([TESSERACT_EXE, "--list-langs"], universal_newlines=True) if not lang_out: # something went wrong raise ValueError("No output from --list-langs") lines = lang_out.splitlines() if len(lines) <= 1: raise ValueError("No output from --list-langs") return sorted(lines[1:]) # ignore first line def extract_text(image_path: str, languages: list[str] = [], verbose: bool = False) -> str: exe_params = [TESSERACT_EXE, image_path, "stdout"] if verbose: exe_params.extend(["-v"]) if languages: exe_params.extend(["-l", "+".join(languages)]) res = subprocess.run(exe_params, capture_output=True, check=True, text=True) if res.stderr: demisto.debug(f"tesseract returned ok but stderr contains warnings: {res.stderr}") return res.stdout def list_languages_command() -> CommandResults: langs = list_languages() return CommandResults( readable_output="## Image OCR Supported Languages\n\n" + "\n".join(["* " + s for s in langs]), raw_response=langs ) def extract_text_command(args: dict, instance_languages: list, skip_corrupted: bool) -> tuple[list, list]: langs = argToList(args.get("langs")) or instance_languages verbose = argToBoolean(args.get("verbose", False)) demisto.debug(f"Using langs settings: {langs}") results, errors = [], [] entry_ids = argToList(args.get("entryid")) for entry_id in entry_ids: try: file_path = demisto.getFilePath(entry_id) if not file_path: raise DemistoException(f"Couldn't find entry id: {entry_id}") if os.path.getsize(file_path["path"]) == 0: demisto.debug(f"File with entry ID {entry_id} is empty, skipping OCR.") results.append( CommandResults( readable_output=f"## Could not process file with entry ID {entry_id} - file is empty", outputs_prefix="File", outputs_key_field="EntryID", outputs={"EntryID": entry_id, "Text": EMPTY_FILE_MSG}, entry_type=EntryType.WARNING, ) ) continue demisto.debug(f"Extracting text from file: {file_path}") res = extract_text(file_path["path"], langs, verbose) file_entry = {"EntryID": entry_id, "Text": res} results.append( CommandResults( readable_output=f"## Image OCR Extracted Text for Entry ID {entry_id}\n\n" + res, outputs_prefix="File", outputs_key_field="EntryID", outputs=file_entry, raw_response=res, ) ) except subprocess.CalledProcessError as cpe: if any(err in cpe.stderr for err in CORRUPTED_ERRORS) and skip_corrupted: file_entry = {"EntryID": entry_id, "Text": CORRUPTED_MSG} results.append( CommandResults( readable_output=f"## Could not process file with entry ID {entry_id} - image is corrupted", outputs_prefix="File", outputs_key_field="EntryID", outputs=file_entry, entry_type=EntryType.WARNING, ) ) else: errors.append( f"An error occurred while trying to process {entry_id=}: " f"Failed {cpe.cmd} execution. Return status: {cpe.returncode}.\n" f"Error:\n{cpe.stderr}\n" f"Stdout:\n{cpe.stdout}" ) except Exception as e: errors.append(f"An error occurred while trying to process {entry_id=}: {e}") return results, errors def run_test_module(instance_languages: list) -> str: try: supported_langs = list_languages() if instance_languages: for language in instance_languages: if language not in supported_langs: raise DemistoException(f"Unsupported language configured: {language}") return "ok" except Exception as exception: raise Exception(f"Failed testing {TESSERACT_EXE}: {exception}") def main() -> None: command = demisto.command() args = demisto.args() params = demisto.params() instance_languages = argToList(params.get("langs")) skip_corrupted = params.get("skip_corrupted") try: if command == "test-module": return_results(run_test_module(instance_languages)) elif command == "image-ocr-list-languages": return_results(list_languages_command()) elif command == "image-ocr-extract-text": results, errors = extract_text_command(args, instance_languages, skip_corrupted) return_results(results) if errors: raise DemistoException("\n".join(errors)) else: raise NotImplementedError(f"Command {command} was not implemented.") except Exception as err: return_error(f"Failed with error(s): {err}") # python2 uses __builtin__ python3 uses builtins if __name__ in ("__builtin__", "builtins", "__main__"): main()