WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

docreader/utils/__init__.py

@@ -0,0 +1,83 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+import os
+import re
+import logging
+
+# 配置日志
+logger = logging.getLogger(__name__)
+
+
+def singleton(cls, *args, **kw):
+    instances = {}
+
+    def _singleton():
+        key = str(cls) + str(os.getpid())
+        if key not in instances:
+            logger.info(f"Creating new singleton instance with key: {key}")
+            instances[key] = cls(*args, **kw)
+        else:
+            logger.info(f"Returning existing singleton instance with key: {key}")
+        return instances[key]
+
+    return _singleton
+
+
+def rmSpace(txt):
+    logger.info(f"Removing spaces from text of length: {len(txt)}")
+    txt = re.sub(r"([^a-z0-9.,\)>]) +([^ ])", r"\1\2", txt, flags=re.IGNORECASE)
+    return re.sub(r"([^ ]) +([^a-z0-9.,\(<])", r"\1\2", txt, flags=re.IGNORECASE)
+
+
+def findMaxDt(fnm):
+    m = "1970-01-01 00:00:00"
+    logger.info(f"Finding maximum date in file: {fnm}")
+    try:
+        with open(fnm, "r") as f:
+            while True:
+                l = f.readline()
+                if not l:
+                    break
+                l = l.strip("\n")
+                if l == "nan":
+                    continue
+                if l > m:
+                    m = l
+        logger.info(f"Maximum date found: {m}")
+    except Exception as e:
+        logger.error(f"Error reading file {fnm} for max date: {str(e)}")
+    return m
+
+
+def findMaxTm(fnm):
+    m = 0
+    logger.info(f"Finding maximum time in file: {fnm}")
+    try:
+        with open(fnm, "r") as f:
+            while True:
+                l = f.readline()
+                if not l:
+                    break
+                l = l.strip("\n")
+                if l == "nan":
+                    continue
+                if int(l) > m:
+                    m = int(l)
+        logger.info(f"Maximum time found: {m}")
+    except Exception as e:
+        logger.error(f"Error reading file {fnm} for max time: {str(e)}")
+    return m

docreader/utils/endecode.py

@@ -0,0 +1,204 @@
+"""
+Encoding and Decoding Utilities Module
+
+This module provides utilities for encoding and decoding various data types,
+with a focus on image and text data conversion:
+- Image encoding/decoding (base64)
+- Text encoding/decoding (multiple character sets)
+- Bytes conversion utilities
+"""
+
+import base64
+import binascii
+import io
+import logging
+from typing import List, Union
+
+import numpy as np
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+
+def decode_image(image: Union[str, bytes, Image.Image, np.ndarray]) -> str:
+    """Convert image to base64 encoded string.
+
+    This function handles multiple image input formats and converts them
+    to a base64 encoded string representation, which is useful for embedding
+    images in JSON, HTML, or other text-based formats.
+
+    Args:
+        image: Image in one of the following formats:
+            - str: File path to an image file
+            - bytes: Raw image bytes data
+            - Image.Image: PIL/Pillow Image object
+            - np.ndarray: NumPy array representing image data
+
+    Returns:
+        str: Base64 encoded string representation of the image
+
+    Raises:
+        ValueError: If the image type is not supported
+
+    Example:
+        >>> # From file path
+        >>> base64_str = decode_image("/path/to/image.png")
+        >>> # From PIL Image
+        >>> from PIL import Image
+        >>> img = Image.open("photo.jpg")
+        >>> base64_str = decode_image(img)
+    """
+    if isinstance(image, str):
+        # Handle file path: read file and encode to base64
+        with open(image, "rb") as image_file:
+            return base64.b64encode(image_file.read()).decode()
+
+    elif isinstance(image, bytes):
+        # Handle raw bytes: directly encode to base64
+        return base64.b64encode(image).decode()
+
+    elif isinstance(image, Image.Image):
+        # Handle PIL Image: save to buffer then encode
+        buffer = io.BytesIO()
+        # Use original format if available, otherwise default to PNG
+        img_format = image.format if image.format else "PNG"
+        image.save(buffer, format=img_format)
+        return base64.b64encode(buffer.getvalue()).decode()
+
+    elif isinstance(image, np.ndarray):
+        # Handle numpy array: convert to PIL Image, then encode as PNG
+        pil_image = Image.fromarray(image)
+        buffer = io.BytesIO()
+        pil_image.save(buffer, format="PNG")
+        return base64.b64encode(buffer.getvalue()).decode()
+
+    raise ValueError(f"Unsupported image type: {type(image)}")
+
+
+def encode_image(image: str, errors="strict") -> bytes:
+    """Decode a base64 encoded image string back to bytes.
+
+    This function converts a base64 encoded string representation of an image
+    back into its original binary bytes format.
+
+    Args:
+        image: Base64 encoded string representation of an image
+        errors: Error handling scheme for decoding errors:
+            - 'strict' (default): Raise binascii.Error on decoding errors
+            - 'ignore': Return empty bytes on decoding errors
+            - Any other name registered with codecs.register_error
+
+    Returns:
+        bytes: Decoded image bytes, or empty bytes if errors='ignore' and decoding fails
+
+    Raises:
+        binascii.Error: If decoding fails and errors='strict'
+
+    Example:
+        >>> base64_str = "iVBORw0KGgoAAAANSUhEUgAAAAUA..."
+        >>> image_bytes = encode_image(base64_str)
+        >>> # With error handling
+        >>> image_bytes = encode_image(base64_str, errors="ignore")
+    """
+    try:
+        # Attempt to decode the base64 string to bytes
+        image_bytes = base64.b64decode(image)
+    except binascii.Error as e:
+        # Handle decoding errors based on the errors parameter
+        if errors == "ignore":
+            return b""
+        else:
+            raise e
+    return image_bytes
+
+
+def encode_bytes(content: str) -> bytes:
+    """Convert a string to bytes using UTF-8 encoding.
+
+    Args:
+        content: String to be encoded
+
+    Returns:
+        bytes: UTF-8 encoded bytes representation of the string
+
+    Example:
+        >>> text = "Hello, 世界"
+        >>> encoded = encode_bytes(text)
+        >>> type(encoded)
+        <class 'bytes'>
+    """
+    return content.encode()
+
+
+def decode_bytes(
+    content: bytes,
+    encodings: List[str] = [
+        "utf-8",
+        "gb18030",
+        "gb2312",
+        "gbk",
+        "big5",
+        "ascii",
+        "latin-1",
+    ],
+) -> str:
+    """Decode bytes to string with automatic encoding detection.
+
+    This function attempts to decode bytes using multiple encoding formats
+    in order of priority. It's particularly useful for handling text files
+    with unknown or mixed encodings, especially for Chinese text.
+
+    The function tries encodings in the provided order and returns the first
+    successful decode. If all encodings fail, it falls back to latin-1 with
+    error replacement to ensure a result is always returned.
+
+    Args:
+        content: Bytes content to be decoded
+        encodings: List of encoding formats to try, in order of priority.
+            Default includes common encodings for Chinese and Western text:
+            - utf-8: Universal encoding (tried first)
+            - gb18030, gb2312, gbk: Chinese encodings (Simplified)
+            - big5: Chinese encoding (Traditional)
+            - ascii, latin-1: Western encodings
+
+    Returns:
+        str: Decoded string content
+
+    Note:
+        - If all encodings fail, latin-1 with error='replace' is used as fallback
+        - The fallback may result in character replacement (�) for invalid bytes
+        - A warning is logged when fallback encoding is used
+
+    Example:
+        >>> # Decode with default encodings
+        >>> text = decode_bytes(b"\\xe4\\xb8\\xad\\xe6\\x96\\x87")  # UTF-8 Chinese
+        >>> print(text)
+        中文
+        >>> # Decode with custom encodings
+        >>> text = decode_bytes(content, encodings=["utf-8", "gbk"])
+    """
+    # Try decoding with each encoding format in order
+    for encoding in encodings:
+        try:
+            text = content.decode(encoding)
+            logger.debug(f"Decode content with {encoding}: {len(text)} characters")
+            return text
+        except UnicodeDecodeError:
+            # This encoding didn't work, try the next one
+            continue
+
+    # Fallback: use latin-1 with error replacement if all encodings fail
+    # latin-1 can decode any byte sequence, but may produce incorrect characters
+    text = content.decode(encoding="latin-1", errors="replace")
+    logger.warning(
+        "Unable to determine correct encoding, using latin-1 as fallback. "
+        "This may cause character issues."
+    )
+    return text
+
+
+if __name__ == "__main__":
+    # Example: Test encode_image with error handling
+    # This demonstrates decoding a base64 string with 'ignore' error mode
+    img = "test![](data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAMgA)test"
+    encode_image(img, errors="ignore")

docreader/utils/request.py

@@ -0,0 +1,149 @@
+import contextlib
+import logging
+import time
+import uuid
+from contextvars import ContextVar
+from logging import LogRecord
+from typing import Optional
+
+# 配置日志
+logger = logging.getLogger(__name__)
+
+# 定义上下文变量
+request_id_var = ContextVar("request_id", default=None)
+_request_start_time_ctx = ContextVar("request_start_time", default=None)
+
+
+def set_request_id(request_id: str) -> None:
+    """设置当前上下文的请求ID"""
+    request_id_var.set(request_id)
+
+
+def get_request_id() -> Optional[str]:
+    """获取当前上下文的请求ID"""
+    return request_id_var.get()
+
+
+class MillisecondFormatter(logging.Formatter):
+    """自定义日志格式化器，只显示毫秒级时间戳(3位数字)而不是微秒(6位)"""
+
+    def formatTime(self, record, datefmt=None):
+        """重写formatTime方法，将微秒格式化为毫秒"""
+        # 先获取标准的格式化时间
+        result = super().formatTime(record, datefmt)
+
+        # 如果使用了包含.%f的格式，则将微秒(6位)截断为毫秒(3位)
+        if datefmt and ".%f" in datefmt:
+            # 格式化的时间字符串应该在最后有6位微秒数
+            parts = result.split(".")
+            if len(parts) > 1 and len(parts[1]) >= 6:
+                # 只保留前3位作为毫秒
+                millis = parts[1][:3]
+                result = f"{parts[0]}.{millis}"
+
+        return result
+
+
+def init_logging_request_id():
+    """
+    Initialize logging to include request ID in log messages.
+    Add the custom filter to all existing handlers
+    """
+    logger.info("Initializing request ID logging")
+    root_logger = logging.getLogger()
+
+    # 添加自定义过滤器到所有处理器
+    for handler in root_logger.handlers:
+        # 添加请求ID过滤器
+        handler.addFilter(RequestIdFilter())
+
+        # 更新格式化器以包含请求ID，调整格式使其更紧凑整齐
+        formatter = logging.Formatter(
+            fmt="%(asctime)s.%(msecs)03d [%(request_id)s] %(levelname)-5s %(name)-20s | %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+        handler.setFormatter(formatter)
+
+    logger.info(
+        f"Updated {len(root_logger.handlers)} handlers with request ID formatting"
+    )
+
+    # 如果没有处理器，添加一个标准输出处理器
+    if not root_logger.handlers:
+        handler = logging.StreamHandler()
+        formatter = logging.Formatter(
+            fmt="%(asctime)s.%(msecs)03d [%(request_id)s] %(levelname)-5s %(name)-20s | %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+        handler.setFormatter(formatter)
+        handler.addFilter(RequestIdFilter())
+        root_logger.addHandler(handler)
+        logger.info("Added new StreamHandler with request ID formatting")
+
+
+class RequestIdFilter(logging.Filter):
+    """Filter that adds request ID to log messages"""
+
+    def filter(self, record: LogRecord) -> bool:
+        request_id = request_id_var.get()
+        if request_id is not None:
+            # 为日志记录添加请求ID属性，使用短格式
+            if len(request_id) > 8:
+                # 截取ID的前8个字符，确保显示整齐
+                short_id = request_id[:8]
+                if "-" in request_id:
+                    # 尝试保留格式，例如 test-req-1-XXX
+                    parts = request_id.split("-")
+                    if len(parts) >= 3:
+                        # 如果格式是 xxx-xxx-n-randompart
+                        short_id = f"{parts[0]}-{parts[1]}-{parts[2]}"
+                record.request_id = short_id
+            else:
+                record.request_id = request_id
+
+            # 添加执行时间属性
+            start_time = _request_start_time_ctx.get()
+            if start_time is not None:
+                elapsed_ms = int((time.time() - start_time) * 1000)
+                record.elapsed_ms = elapsed_ms
+                # 添加执行时间到消息中
+                if not hasattr(record, "message_with_elapsed"):
+                    record.message_with_elapsed = True
+                    record.msg = f"{record.msg} (elapsed: {elapsed_ms}ms)"
+        else:
+            # 如果没有请求ID，使用占位符
+            record.request_id = "no-req-id"
+
+        return True
+
+
+@contextlib.contextmanager
+def request_id_context(request_id: str = None):
+    """Context manager that sets a request ID for the current context
+
+    Args:
+        request_id: 要使用的请求ID，如果为None则自动生成
+
+    Example:
+        with request_id_context("req-123"):
+            # 在这个代码块中的所有日志都会包含请求ID req-123
+            logging.info("Processing request")
+    """
+    # Generate or use provided request ID
+    req_id = request_id or str(uuid.uuid4())
+
+    # Set start time and request ID
+    start_time = time.time()
+    req_token = request_id_var.set(req_id)
+    time_token = _request_start_time_ctx.set(start_time)
+
+    logger.info(f"Starting new request with ID: {req_id}")
+
+    try:
+        yield request_id_var.get()
+    finally:
+        # Log completion and reset context vars
+        elapsed_ms = int((time.time() - start_time) * 1000)
+        logger.info(f"Request {req_id} completed in {elapsed_ms}ms")
+        request_id_var.reset(req_token)
+        _request_start_time_ctx.reset(time_token)

docreader/utils/split.py

@@ -0,0 +1,80 @@
+import re
+from typing import Callable, List
+
+
+def split_text_keep_separator(text: str, separator: str) -> List[str]:
+    """Split text with separator and keep the separator at the end of each split.
+    
+    Args:
+        text: The input text to split
+        separator: The separator string to split by
+        
+    Returns:
+        List of text chunks with separator preserved at the start of each chunk (except first)
+        
+    Example:
+        >>> split_text_keep_separator("Hello\nWorld\nTest", "\n")
+        ["Hello", "\nWorld", "\nTest"]
+    """
+    # Split text by separator
+    parts = text.split(separator)
+    # Add separator back to the beginning of each part (except the first one)
+    result = [separator + s if i > 0 else s for i, s in enumerate(parts)]
+    # Filter out empty strings
+    return [s for s in result if s]
+
+
+def split_by_sep(sep: str, keep_sep: bool = True) -> Callable[[str], List[str]]:
+    """Create a function that splits text by a given separator.
+    
+    Args:
+        sep: The separator string to split by
+        keep_sep: If True, keep the separator in the result; if False, discard it
+        
+    Returns:
+        A callable function that takes text and returns a list of split strings
+    """
+    if keep_sep:
+        return lambda text: split_text_keep_separator(text, sep)
+    else:
+        return lambda text: text.split(sep)
+
+
+def split_by_char() -> Callable[[str], List[str]]:
+    """Create a function that splits text into individual characters.
+    
+    Returns:
+        A callable function that takes text and returns a list of characters
+    """
+    return lambda text: list(text)
+
+
+def split_by_regex(regex: str) -> Callable[[str], List[str]]:
+    """Create a function that splits text by a regex pattern.
+    
+    Args:
+        regex: The regular expression pattern to split by
+        
+    Returns:
+        A callable function that takes text and returns a list of split strings
+        The regex pattern is captured, so the separators are included in the result
+    """
+    # Compile regex with capturing group to keep separators in result
+    pattern = re.compile(f"({regex})")
+    # Split by pattern and filter out None/empty values
+    return lambda text: list(filter(None, pattern.split(text)))
+
+
+def match_by_regex(regex: str) -> Callable[[str], bool]:
+    """Create a function that checks if text matches a regex pattern.
+    
+    Args:
+        regex: The regular expression pattern to match against
+        
+    Returns:
+        A callable function that takes text and returns True if it matches the pattern
+    """
+    # Compile the regex pattern for efficient reuse
+    pattern = re.compile(regex)
+    # Return a function that checks if text matches the pattern from the start
+    return lambda text: bool(pattern.match(text))

docreader/utils/ssrf.py

@@ -0,0 +1,285 @@
+"""SSRF URL validation for docreader outbound HTTP requests.
+
+Mirrors the core policy in internal/utils/security.go so redirect targets
+during Playwright navigation are blocked the same way as Go-side imports.
+"""
+
+from __future__ import annotations
+
+import ipaddress
+import os
+import re
+import socket
+from functools import lru_cache
+from typing import FrozenSet, Optional, Tuple, Union
+from urllib.parse import urlparse
+
+RESTRICTED_HOSTNAMES: FrozenSet[str] = frozenset(
+    {
+        "localhost",
+        "127.0.0.1",
+        "::1",
+        "0.0.0.0",
+        "metadata.google.internal",
+        "metadata.tencentyun.com",
+        "metadata.aws.internal",
+        "host.docker.internal",
+        "gateway.docker.internal",
+        "kubernetes.docker.internal",
+        "kubernetes",
+        "kubernetes.default",
+        "kubernetes.default.svc",
+        "kubernetes.default.svc.cluster.local",
+    }
+)
+
+RESTRICTED_SUFFIXES: Tuple[str, ...] = (
+    ".local",
+    ".localhost",
+    ".internal",
+    ".corp",
+    ".lan",
+    ".home",
+    ".localdomain",
+    ".svc.cluster.local",
+    ".pod.cluster.local",
+)
+
+EXTRA_RESTRICTED_CIDRS: Tuple[Union[ipaddress.IPv4Network, ipaddress.IPv6Network], ...] = tuple(
+    ipaddress.ip_network(cidr)
+    for cidr in (
+        "100.64.0.0/10",
+        "198.18.0.0/15",
+        "198.51.100.0/24",
+        "203.0.113.0/24",
+        "192.0.0.0/24",
+        "192.0.2.0/24",
+        "0.0.0.0/8",
+        "240.0.0.0/4",
+        "255.255.255.255/32",
+        "172.17.0.0/16",
+        "172.18.0.0/16",
+        "172.19.0.0/16",
+        "172.20.0.0/16",
+    )
+)
+
+BLOCKED_PORTS: FrozenSet[str] = frozenset(
+    {
+        "22",
+        "23",
+        "25",
+        "445",
+        "3389",
+        "5432",
+        "3306",
+        "6379",
+        "27017",
+        "9200",
+        "2379",
+        "2380",
+        "8500",
+        "4001",
+    }
+)
+
+_IP_LIKE_PATTERNS = (
+    re.compile(r"^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$"),
+    re.compile(r"^\d{8,10}$"),
+    re.compile(r"^0[0-7]+\."),
+    re.compile(r"(?i)^0x[0-9a-f]+\."),
+    re.compile(r"(?i)^0x[0-9a-f]{6,8}$"),
+    re.compile(r"(?i)^[0-9a-f:]+::[0-9a-f:]*$"),
+    re.compile(r"(?i)^[0-9a-f]{1,4}(:[0-9a-f]{1,4}){7}$"),
+)
+
+
+def _normalize_url(raw_url: str) -> str:
+    if "://" not in raw_url:
+        return f"https://{raw_url}"
+    return raw_url
+
+
+@lru_cache(maxsize=1)
+def _load_whitelist() -> Tuple[FrozenSet[str], Tuple[str, ...], Tuple[Union[ipaddress.IPv4Network, ipaddress.IPv6Network], ...]]:
+    entries: list[str] = []
+    for env_key in ("SSRF_WHITELIST", "SSRF_WHITELIST_EXTRA"):
+        raw = os.environ.get(env_key, "")
+        if raw.strip():
+            entries.extend(part.strip() for part in raw.split(",") if part.strip())
+
+    exact_hosts: set[str] = set()
+    suffix_hosts: list[str] = []
+    cidr_nets: list[Union[ipaddress.IPv4Network, ipaddress.IPv6Network]] = []
+    for entry in entries:
+        lowered = entry.lower()
+        if lowered.startswith("*."):
+            suffix_hosts.append(lowered[1:])
+            continue
+        if "/" in lowered:
+            try:
+                cidr_nets.append(ipaddress.ip_network(lowered, strict=False))
+            except ValueError:
+                continue
+            continue
+        exact_hosts.add(lowered)
+    return frozenset(exact_hosts), tuple(suffix_hosts), tuple(cidr_nets)
+
+
+def _is_whitelisted(hostname: str) -> bool:
+    lowered = hostname.lower()
+    exact_hosts, suffix_hosts, cidr_nets = _load_whitelist()
+    if lowered in exact_hosts:
+        return True
+    for suffix in suffix_hosts:
+        if lowered.endswith(suffix) or lowered == suffix.lstrip("."):
+            return True
+    try:
+        ip = ipaddress.ip_address(lowered)
+    except ValueError:
+        return False
+    return any(ip in net for net in cidr_nets)
+
+
+def _is_ip_like_hostname(hostname: str) -> bool:
+    return any(pattern.search(hostname) for pattern in _IP_LIKE_PATTERNS)
+
+
+def _is_restricted_ip(ip: Union[ipaddress.IPv4Address, ipaddress.IPv6Address]) -> Optional[str]:
+    if ip.is_private:
+        return "private IP address"
+    if ip.is_loopback:
+        return "loopback address"
+    if ip.is_link_local:
+        return "link-local address"
+    if ip.is_multicast:
+        return "multicast address"
+    if ip.is_unspecified:
+        return "unspecified address"
+    if isinstance(ip, ipaddress.IPv4Address):
+        for net in EXTRA_RESTRICTED_CIDRS:
+            if isinstance(net, ipaddress.IPv4Network) and ip in net:
+                return f"restricted range {net}"
+    if isinstance(ip, ipaddress.IPv6Address):
+        embedded_ipv4 = ip.ipv4_mapped
+        if embedded_ipv4 is not None:
+            reason = _is_restricted_ip(embedded_ipv4)
+            if reason:
+                return f"IPv4-mapped {reason}"
+        if ip.sixtofour is not None:
+            reason = _is_restricted_ip(ip.sixtofour)
+            if reason:
+                return f"6to4-embedded {reason}"
+        if ip.teredo is not None:
+            server_ip, client_ip = ip.teredo
+            for label, embedded_ip in (("server", server_ip), ("client", client_ip)):
+                reason = _is_restricted_ip(embedded_ip)
+                if reason:
+                    return f"Teredo {label} embeds {reason}"
+        # NAT64's well-known prefix (64:ff9b::/96, RFC 6052) and the deprecated
+        # IPv4-compatible form (::a.b.c.d) also carry a plain IPv4 address, but
+        # ipaddress exposes no accessor for either, so is_private and
+        # .sixtofour never see the payload. Without these two branches
+        # ::169.254.169.254 reads as an ordinary public IPv6 address.
+        # Mirrors internal/ipclass on the Go side.
+        packed = ip.packed
+        if packed[0:4] == b"\x00\x64\xff\x9b" and packed[4:12] == bytes(8):
+            reason = _is_restricted_ip(ipaddress.IPv4Address(packed[12:16]))
+            if reason:
+                return f"NAT64-embedded {reason}"
+        if packed[0:12] == bytes(12):
+            # :: and ::1 already returned above; ::ffff:a.b.c.d is handled by
+            # the ipv4_mapped branch.
+            reason = _is_restricted_ip(ipaddress.IPv4Address(packed[12:16]))
+            if reason:
+                return f"IPv4-compatible {reason}"
+        # Site-local (fec0::/10)
+        if (ip.packed[0] == 0xFE) and (ip.packed[1] & 0xC0) == 0xC0:
+            return "site-local IPv6 address"
+    return None
+
+
+def _resolve_host_ips(hostname: str) -> Tuple[Tuple[Union[ipaddress.IPv4Address, ipaddress.IPv6Address], ...], Optional[str]]:
+    try:
+        infos = socket.getaddrinfo(hostname, None, type=socket.SOCK_STREAM)
+    except socket.gaierror as exc:
+        return (), f"DNS resolution failed for hostname {hostname}: {exc}"
+    ips: list[Union[ipaddress.IPv4Address, ipaddress.IPv6Address]] = []
+    seen: set[str] = set()
+    for info in infos:
+        sockaddr = info[4]
+        if not sockaddr:
+            continue
+        ip_str = sockaddr[0]
+        if ip_str in seen:
+            continue
+        seen.add(ip_str)
+        try:
+            ips.append(ipaddress.ip_address(ip_str))
+        except ValueError:
+            continue
+    if not ips:
+        return (), f"DNS resolution failed for hostname {hostname}: no addresses"
+    return tuple(ips), None
+
+
+def is_ssrf_safe_url(raw_url: str) -> Tuple[bool, str]:
+    """Return (safe, reason). reason is empty when safe is True."""
+    if not raw_url or not raw_url.strip():
+        return False, "URL is empty"
+
+    normalized = _normalize_url(raw_url.strip())
+    parsed = urlparse(normalized)
+    scheme = (parsed.scheme or "").lower()
+    if scheme not in {"http", "https"}:
+        return False, f"invalid scheme: {scheme or '(none)'} (only http/https allowed)"
+
+    hostname = (parsed.hostname or "").strip()
+    if not hostname:
+        return False, "URL has no hostname"
+
+    hostname_lower = hostname.lower()
+    if _is_whitelisted(hostname_lower):
+        return True, ""
+
+    if hostname_lower in RESTRICTED_HOSTNAMES:
+        return False, f"hostname {hostname_lower} is restricted"
+
+    for suffix in RESTRICTED_SUFFIXES:
+        if hostname_lower.endswith(suffix):
+            return False, f"hostname suffix {suffix} is restricted"
+
+    try:
+        ipaddress.ip_address(hostname_lower)
+        return False, "direct IP address access is not allowed, use domain name or add to SSRF_WHITELIST"
+    except ValueError:
+        pass
+
+    if _is_ip_like_hostname(hostname_lower):
+        return False, "IP-like hostname format is not allowed"
+
+    resolved_ips, resolve_err = _resolve_host_ips(hostname_lower)
+    if resolve_err:
+        return False, resolve_err
+
+    for resolved_ip in resolved_ips:
+        reason = _is_restricted_ip(resolved_ip)
+        if reason:
+            return (
+                False,
+                f"hostname {hostname_lower} resolves to restricted IP {resolved_ip}: {reason}",
+            )
+
+    try:
+        port = parsed.port
+    except ValueError as exc:
+        return False, f"invalid port: {exc}"
+    if port is not None and str(port) in BLOCKED_PORTS:
+        return False, f"port {port} is blocked for security reasons"
+
+    return True, ""
+
+
+def reset_ssrf_whitelist_cache_for_test() -> None:
+    """Clear cached whitelist entries (for unit tests only)."""
+    _load_whitelist.cache_clear()

docreader/utils/tempfile.py

@@ -0,0 +1,77 @@
+import logging
+import os
+import tempfile
+
+logger = logging.getLogger(__name__)
+
+
+class TempFileContext:
+    def __init__(self, file_content: bytes, suffix: str):
+        """
+        Initialize the context
+        :param file_content: Byte data to write to file
+        :param suffix: File suffix
+        """
+        self.file_content = file_content
+        self.suffix = suffix
+        self.file = None
+
+    def __enter__(self):
+        """
+        Create file when entering context
+        """
+        self.temp_file = tempfile.NamedTemporaryFile(suffix=self.suffix, delete=False)
+        self.temp_file.write(self.file_content)
+        self.temp_file.flush()
+        logger.info(
+            f"Saved {self.suffix} content to temporary file: {self.temp_file.name}"
+        )
+        return self.temp_file.name
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Delete file when exiting context
+        """
+        if self.temp_file:
+            self.temp_file.close()
+            if os.path.exists(self.temp_file.name):
+                os.remove(self.temp_file.name)
+            logger.info(f"File {self.temp_file.name} has been deleted.")
+        # Return False to propagate exception (if any exception occurred)
+        return False
+
+
+class TempDirContext:
+    def __init__(self):
+        """
+        Initialize the context
+        """
+        self.temp_dir = None
+
+    def __enter__(self):
+        """
+        Create directory when entering context
+        """
+        self.temp_dir = tempfile.TemporaryDirectory()
+        logger.info(f"Created temporary directory: {self.temp_dir.name}")
+        return self.temp_dir.name
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Delete directory when exiting context
+        """
+        if self.temp_dir and os.path.exists(self.temp_dir.name):
+            self.temp_dir.cleanup()
+            logger.info(f"Directory {self.temp_dir.name} has been deleted.")
+        # Return False to propagate exception (if any exception occurred)
+        return False
+
+
+if __name__ == "__main__":
+    example_bytes = b"Hello, this is a test file."
+    file_name = "test_file.txt"
+
+    # Using with statement
+    with TempFileContext(example_bytes, file_name) as temp_file:
+        # File operations can be performed within the context
+        print(f"Does file {file_name} exist: {os.path.exists(file_name)}")