Repository: barvhaim/AutoC
Pull request #2: added ocr capabilities to html_extractor
Author: login: netfried
Created: 2025-04-09T09:06:55Z | Merged: 2025-04-22T07:45:58Z

remove duplicate IOCs


Diff:
diff --git a/.env.sample b/.env.sample
index e8f0227..a6525d2 100644
--- a/.env.sample
+++ b/.env.sample
@@ -18,4 +18,5 @@ LANGSMITH_PROJECT=
 
 VIRUSTOTAL_API_KEY=
 
-SKIP_QNA=false
\ No newline at end of file
+SKIP_QNA=false
+ANALYZE_BLOG_IMAGES=false
\ No newline at end of file
diff --git a/backend/extractors/iocs_extractor.py b/backend/extractors/iocs_extractor.py
index ad58a15..1ef52dd 100644
--- a/backend/extractors/iocs_extractor.py
+++ b/backend/extractors/iocs_extractor.py
@@ -68,7 +68,7 @@ def _extract_ioc(self, ioc_type: IOCType) -> RunnableSequence:
         )
 
     def extract_iocs_from_text(self) -> List[IOC]:
-        iocs = []
+        iocs, val_l = [], []
         tasks = {
             IOCType.URL.name: self._extract_ioc(IOCType.URL),
             IOCType.IP.name: self._extract_ioc(IOCType.IP),
@@ -77,6 +77,11 @@ def extract_iocs_from_text(self) -> List[IOC]:
             IOCType.CHROME_EXTENSION.name: self._extract_ioc(IOCType.CHROME_EXTENSION),
         }
         res = RunnableParallel(**tasks).invoke(input={})
-        for iocs_of_type in res.values():
-            iocs.extend(iocs_of_type)
+
+        # make sure we remove duplicate iocs
+        for ioc_of_type in res.values():
+            for ioc in ioc_of_type:
+                if ioc.value not in val_l:
+                    iocs.append(ioc)
+                    val_l.append(ioc.value)
         return iocs
diff --git a/backend/parsers/html_parser.py b/backend/parsers/html_parser.py
index f9c6d75..cc1a1a2 100644
--- a/backend/parsers/html_parser.py
+++ b/backend/parsers/html_parser.py
@@ -1,24 +1,97 @@
+import base64
 import multiprocessing
 from typing import Optional
+from itertools import chain
+from PIL import Image
+import pytesseract
+from io import BytesIO
 import logging
 import re
 import docling.exceptions
 from docling.document_converter import DocumentConverter
 import requests
 from bs4 import BeautifulSoup
+from urllib.parse import urljoin
+
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
 class HTMLParser:
-    def __init__(self, url: str):
+    def __init__(self, url: str, use_ocr: bool = False):
         self.url = url
         self.converter = DocumentConverter()
+        self.use_ocr: bool = use_ocr
         self.headers = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
         }
 
+    def extract_text_from_image(self, image_src: str) -> str:
+        img = None
+
+        # Load image from URL
+        if image_src.startswith('http'):
+            response = requests.get(image_src)
+            if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
+                img = Image.open(BytesIO(response.content))
+            else:
+                content_type = response.headers.get('Content-Type')
+                logger.info(
+                    f"Skipped: {image_src} (status: {response.status_code}, content-type: {content_type})")
+
+        # Load from base64 (optional enhancement)
+        elif image_src.startswith('data:image'):
+            base64_data = image_src.split(",")[1]
+            img = Image.open(BytesIO(base64.b64decode(base64_data)))
+
+        elif len(self._overlap_with_base_url(image_src)) > 0:
+            response = requests.get(urljoin(self.url, image_src))
+            img = Image.open(BytesIO(response.content))
+        else:
+            logger.info(f"Unsupported image src format: {image_src}")
+
+        # OCR on the image
+        text = pytesseract.image_to_string(img)
+        return text
+
+    def _overlap_with_base_url(self, img_url: str) -> str:
+        max_len = min(len(self.url), len(img_url))
+        for i in range(max_len, 0, -1):
+            if self.url[-i:] == img_url[:i]:
+                return self.url[-i:]
+        return ""
+
+    @staticmethod
+    def _extract_image_src(picture_tag, html_type: str) -> str:
+        if html_type == 'picture_tag':
+            source_tag = picture_tag.find('source')
+            if source_tag and source_tag.get('srcset'):
+                return source_tag.get('srcset').split(' ')[0]
+        else:
+            return picture_tag.get('src')
+
+    def _extract_text_from_images_in_html(self, soup_obj: BeautifulSoup) -> str:
+        res = []
+        img_tags = soup_obj.find_all('img')
+        picture_tags = soup_obj.find_all('picture')
+        combined_tags = chain(
+            zip(picture_tags, ['picture_tag'] * len(picture_tags)),
+            zip(img_tags, ['image_tag'] * len(img_tags))
+        )
+        for i, (tag, html_type) in enumerate(combined_tags):
+            src = self._extract_image_src(tag, html_type)
+
+            if not src:
+                continue
+
+            try:
+                text = self.extract_text_from_image(image_src=src)
+                res.append(text)
+            except Exception as e:
+                logger.info(f"Failed to process image {i} ({src}): {e}")
+        return ''.join(res)
+
     @staticmethod
     def _extract_article_content_from_markdown(markdown: str) -> str:
         # Regex to match the title and everything after it, TODO: improve method (?)
@@ -63,6 +136,18 @@ def _extract_textual_content_with_docling(
             logger.warning("Failed to extract blog content using docling")
         result_queue.put(None)
 
+    def _fetch_ocr_data_with_beautifulsoup(self) -> Optional[str]:
+        logger.info("Attempting to extract text from images using BeautifulSoup")
+        try:
+            response = requests.get(self.url, headers=self.headers)
+            response.raise_for_status()
+            soup = BeautifulSoup(response.content, "html.parser")
+
+            image_text_str = self._extract_text_from_images_in_html(soup)
+            return image_text_str.strip()
+        except Exception as e:
+            logger.warning(f"Failed to extract text from images using BeautifulSoup: {e}")
+
     def _extract_textual_content_with_beautifulsoup(self) -> Optional[str]:
         logger.info("Attempting to extract HTML data using BeautifulSoup")
         try:
@@ -73,25 +158,30 @@ def _extract_textual_content_with_beautifulsoup(self) -> Optional[str]:
         return None
 
     def get_textual_content(self) -> Optional[str]:
-        result_queue = multiprocessing.Queue()
-        process = multiprocessing.Process(
-            target=self._extract_textual_content_with_docling, args=(result_queue,)
-        )
-        process.start()
-        process.join(timeout=60)
-
-        if process.is_alive():
-            logger.warning(
-                "Docling extraction is taking too long, falling back to BeautifulSoup"
+        use_multi_proccesing = True
+        if use_multi_proccesing:
+            result_queue = multiprocessing.Queue(maxsize=1)
+            process = multiprocessing.Process(
+                target=self._extract_textual_content_with_docling, args=(result_queue,)
             )
-            process.terminate()
-            process.join()
-
-        content = result_queue.get() if not result_queue.empty() else None
-
+            process.start()
+            process.join(timeout=60)
+
+            if process.is_alive():
+                logger.warning(
+                    "Docling extraction is taking too long, falling back to BeautifulSoup"
+                )
+                process.terminate()
+                process.join()
+
+            content = result_queue.get() if not result_queue.empty() else None
+        else:
+            content = None
         if content is None:
             content = self._extract_textual_content_with_beautifulsoup()
-
+        if self.use_ocr:
+            ocr_text = self._fetch_ocr_data_with_beautifulsoup()
+            content = content + ocr_text if content else ocr_text
         if content:
             return content
 
diff --git a/backend/pipeline/nodes.py b/backend/pipeline/nodes.py
index ec9eede..53a1e5b 100644
--- a/backend/pipeline/nodes.py
+++ b/backend/pipeline/nodes.py
@@ -26,7 +26,7 @@ def html_extractor_node(state: PipelineState) -> Command:
         return Command(goto=END, update={"error": "No blog URL provided"})
 
     logger.info(f"Extracting content from {url}")
-    parser = HTMLParser(url=url)
+    parser = HTMLParser(url=url, use_ocr=os.getenv("ANALYZE_BLOG_IMAGES", 'false') == 'true')
     article_textual_content = parser.get_textual_content()
 
     return Command(
diff --git a/requirements.txt b/requirements.txt
index 5725797..51f8b49 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,6 +14,7 @@ requests
 beautifulsoup4
 pyyaml
 websockets
+pytesseract
 
 # consumption
 fastapi


Reviews:
user: login: freunda
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T11:44:43Z
user: login: netfried
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T07:48:53Z
user: login: netfried
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T07:49:01Z
user: login: netfried
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-10T07:49:05Z
user: login: freunda
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-20T11:50:21Z

Review comments:
diff_hunk: @@ -68,7 +68,7 @@ def _extract_ioc(self, ioc_type: IOCType) -> RunnableSequence:
         )
 
     def extract_iocs_from_text(self) -> List[IOC]:
-        iocs = []
+        iocs, val_l = [], []
path: backend/extractors/iocs_extractor.py
user: login: freunda
body: You can use set here to avoid duplications
```
def extract_iocs_from_text(self) -> List[IOC]:
    iocs = set()  # Initialize as a set
    tasks = {
        IOCType.URL.name: self._extract_ioc(IOCType.URL),
        IOCType.IP.name: self._extract_ioc(IOCType.IP),
        IOCType.MD5.name: self._extract_ioc(IOCType.MD5),
        IOCType.SHA256.name: self._extract_ioc(IOCType.SHA256),
        IOCType.CHROME_EXTENSION.name: self._extract_ioc(IOCType.CHROME_EXTENSION),
    }
    res = RunnableParallel(**tasks).invoke(input={})
    for iocs_of_type in res.values():
        iocs.update(iocs_of_type)  # Use update to add elements to the set
    return list(iocs)  # Convert the set back to a list
```
created_at: 2025-04-09T11:24:00Z
author_association: COLLABORATOR
diff_hunk: @@ -12,13 +17,80 @@
 
 
 class HTMLParser:
-    def __init__(self, url: str):
+    def __init__(self, url: str, use_ocr: bool = True):
         self.url = url
         self.converter = DocumentConverter()
+        self.use_ocr: bool = use_ocr
         self.headers = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
         }
 
+    def extract_text_from_image(self, image_src) -> str:
+        img = None
+
+        # Load image from URL
+        if image_src.startswith('http'):
+            response = requests.get(image_src)
+            if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
+                img = Image.open(BytesIO(response.content))
+            else:
+                content_type = response.headers.get('Content-Type')
+                print(
path: backend/parsers/html_parser.py
user: login: freunda
body: Use logger instead of print
created_at: 2025-04-09T11:29:05Z
author_association: COLLABORATOR
diff_hunk: @@ -12,13 +17,80 @@
 
 
 class HTMLParser:
-    def __init__(self, url: str):
+    def __init__(self, url: str, use_ocr: bool = True):
         self.url = url
         self.converter = DocumentConverter()
+        self.use_ocr: bool = use_ocr
         self.headers = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
         }
 
+    def extract_text_from_image(self, image_src) -> str:
+        img = None
+
+        # Load image from URL
+        if image_src.startswith('http'):
+            response = requests.get(image_src)
+            if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
+                img = Image.open(BytesIO(response.content))
+            else:
+                content_type = response.headers.get('Content-Type')
+                print(
+                    f"Skipped: {image_src} (status: {response.status_code}, content-type: {content_type})")
+
+        # Load from base64 (optional enhancement)
+        elif image_src.startswith('data:image'):
+            base64_data = image_src.split(",")[1]
+            img = Image.open(BytesIO(base64.b64decode(base64_data)))
+
+        elif len(self._overlap_with_base_url(image_src)) > 0:
+            overlap = self._overlap_with_base_url(image_src)
path: backend/parsers/html_parser.py
user: login: freunda
body: You better use urljoin from urllib.parse, as it deals with edge cases
created_at: 2025-04-09T11:39:45Z
author_association: COLLABORATOR
diff_hunk: @@ -68,7 +68,7 @@ def _extract_ioc(self, ioc_type: IOCType) -> RunnableSequence:
         )
 
     def extract_iocs_from_text(self) -> List[IOC]:
-        iocs = []
+        iocs, val_l = [], []
path: backend/extractors/iocs_extractor.py
user: login: netfried
body: iocs_of_type is a list of <class 'backend.data_model.ioc.IOC'> items.
if we want to use set we need to keep the nested loop without the if statement.
unfortunately set.update() destroy the class <class 'backend.data_model.ioc.IOC'> and split between the type and the value.
I am keeping the current solution.
created_at: 2025-04-10T07:48:53Z
author_association: COLLABORATOR
diff_hunk: @@ -12,13 +17,80 @@
 
 
 class HTMLParser:
-    def __init__(self, url: str):
+    def __init__(self, url: str, use_ocr: bool = True):
         self.url = url
         self.converter = DocumentConverter()
+        self.use_ocr: bool = use_ocr
         self.headers = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
         }
 
+    def extract_text_from_image(self, image_src) -> str:
+        img = None
+
+        # Load image from URL
+        if image_src.startswith('http'):
+            response = requests.get(image_src)
+            if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
+                img = Image.open(BytesIO(response.content))
+            else:
+                content_type = response.headers.get('Content-Type')
+                print(
path: backend/parsers/html_parser.py
user: login: netfried
body: done
created_at: 2025-04-10T07:49:00Z
author_association: COLLABORATOR
diff_hunk: @@ -12,13 +17,80 @@
 
 
 class HTMLParser:
-    def __init__(self, url: str):
+    def __init__(self, url: str, use_ocr: bool = True):
         self.url = url
         self.converter = DocumentConverter()
+        self.use_ocr: bool = use_ocr
         self.headers = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
         }
 
+    def extract_text_from_image(self, image_src) -> str:
+        img = None
+
+        # Load image from URL
+        if image_src.startswith('http'):
+            response = requests.get(image_src)
+            if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''):
+                img = Image.open(BytesIO(response.content))
+            else:
+                content_type = response.headers.get('Content-Type')
+                print(
+                    f"Skipped: {image_src} (status: {response.status_code}, content-type: {content_type})")
+
+        # Load from base64 (optional enhancement)
+        elif image_src.startswith('data:image'):
+            base64_data = image_src.split(",")[1]
+            img = Image.open(BytesIO(base64.b64decode(base64_data)))
+
+        elif len(self._overlap_with_base_url(image_src)) > 0:
+            overlap = self._overlap_with_base_url(image_src)
path: backend/parsers/html_parser.py
user: login: netfried
body: done
created_at: 2025-04-10T07:49:05Z
author_association: COLLABORATOR

PR comments:
user: login: barvhaim
created_at: 2025-04-09T19:26:14Z
body: @netfried update the requirements.txt file with required libs for example `pytesseract`
author_association: OWNER
user: login: netfried
created_at: 2025-04-10T07:49:33Z
body: updated the requirements.txt with pytesseract
author_association: COLLABORATOR