""" Watermark Service for MED-009: Dynamic Watermark for Downloads This service provides functions to add watermarks to image and PDF files containing user information for audit and tracking purposes. Watermark content includes: - User name - Employee ID (or email as fallback) - Download timestamp """ import io import logging import math from datetime import datetime from typing import Optional, Tuple import fitz # PyMuPDF from PIL import Image, ImageDraw, ImageFont logger = logging.getLogger(__name__) class WatermarkService: """Service for adding watermarks to downloaded files.""" # Watermark configuration WATERMARK_OPACITY = 0.3 # 30% opacity for semi-transparency WATERMARK_ANGLE = -45 # Diagonal angle in degrees WATERMARK_FONT_SIZE = 24 WATERMARK_COLOR = (128, 128, 128) # Gray color for watermark WATERMARK_SPACING = 200 # Spacing between repeated watermarks @staticmethod def _format_watermark_text( user_name: str, employee_id: Optional[str] = None, download_time: Optional[datetime] = None ) -> str: """ Format the watermark text with user information. Args: user_name: Name of the user employee_id: Employee ID (工號) - uses 'N/A' if not provided download_time: Time of download (defaults to now) Returns: Formatted watermark text """ if download_time is None: download_time = datetime.now() time_str = download_time.strftime("%Y-%m-%d %H:%M:%S") emp_id = employee_id if employee_id else "N/A" return f"{user_name} ({emp_id}) - {time_str}" @staticmethod def _get_font(size: int = 24) -> ImageFont.FreeTypeFont: """Get a font for the watermark. Falls back to default if custom font not available.""" try: # Try to use a common system font (macOS) return ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", size) except (OSError, IOError): try: # Try Linux font return ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", size) except (OSError, IOError): try: # Try Windows font return ImageFont.truetype("C:/Windows/Fonts/arial.ttf", size) except (OSError, IOError): # Fall back to default bitmap font return ImageFont.load_default() def add_image_watermark( self, image_bytes: bytes, user_name: str, employee_id: Optional[str] = None, download_time: Optional[datetime] = None ) -> Tuple[bytes, str]: """ Add a semi-transparent diagonal watermark to an image. Args: image_bytes: The original image as bytes user_name: Name of the user downloading the file employee_id: Employee ID of the user (工號) download_time: Time of download (defaults to now) Returns: Tuple of (watermarked image bytes, output format) Raises: Exception: If watermarking fails """ # Open the image original = Image.open(io.BytesIO(image_bytes)) # Convert to RGBA if necessary for transparency support if original.mode != 'RGBA': image = original.convert('RGBA') else: image = original.copy() # Create a transparent overlay for the watermark watermark_layer = Image.new('RGBA', image.size, (255, 255, 255, 0)) draw = ImageDraw.Draw(watermark_layer) # Get watermark text watermark_text = self._format_watermark_text(user_name, employee_id, download_time) # Get font font = self._get_font(self.WATERMARK_FONT_SIZE) # Calculate text size bbox = draw.textbbox((0, 0), watermark_text, font=font) text_width = bbox[2] - bbox[0] text_height = bbox[3] - bbox[1] # Create a larger canvas for the rotated text pattern diagonal = int(math.sqrt(image.size[0]**2 + image.size[1]**2)) pattern_size = diagonal * 2 # Create pattern layer pattern = Image.new('RGBA', (pattern_size, pattern_size), (255, 255, 255, 0)) pattern_draw = ImageDraw.Draw(pattern) # Draw repeated watermark text across the pattern opacity = int(255 * self.WATERMARK_OPACITY) watermark_color = (*self.WATERMARK_COLOR, opacity) y = 0 row = 0 while y < pattern_size: x = -text_width if row % 2 else 0 # Offset alternate rows while x < pattern_size: pattern_draw.text((x, y), watermark_text, font=font, fill=watermark_color) x += text_width + self.WATERMARK_SPACING y += text_height + self.WATERMARK_SPACING row += 1 # Rotate the pattern rotated_pattern = pattern.rotate( self.WATERMARK_ANGLE, expand=False, center=(pattern_size // 2, pattern_size // 2) ) # Crop to original image size (centered) crop_x = (pattern_size - image.size[0]) // 2 crop_y = (pattern_size - image.size[1]) // 2 cropped_pattern = rotated_pattern.crop(( crop_x, crop_y, crop_x + image.size[0], crop_y + image.size[1] )) # Composite the watermark onto the image watermarked = Image.alpha_composite(image, cropped_pattern) # Determine output format original_format = original.format or 'PNG' if original_format.upper() == 'JPEG': # Convert back to RGB for JPEG (no alpha channel) watermarked = watermarked.convert('RGB') output_format = 'JPEG' else: output_format = 'PNG' # Save to bytes output = io.BytesIO() watermarked.save(output, format=output_format, quality=95) output.seek(0) logger.info( f"Image watermark applied successfully for user {user_name} " f"(employee_id: {employee_id})" ) return output.getvalue(), output_format.lower() def add_pdf_watermark( self, pdf_bytes: bytes, user_name: str, employee_id: Optional[str] = None, download_time: Optional[datetime] = None ) -> bytes: """ Add a semi-transparent diagonal watermark to a PDF using PyMuPDF. Args: pdf_bytes: The original PDF as bytes user_name: Name of the user downloading the file employee_id: Employee ID of the user (工號) download_time: Time of download (defaults to now) Returns: Watermarked PDF as bytes Raises: Exception: If watermarking fails """ # Get watermark text watermark_text = self._format_watermark_text(user_name, employee_id, download_time) # Open the PDF with PyMuPDF doc = fitz.open(stream=pdf_bytes, filetype="pdf") page_count = len(doc) # Process each page for page_num in range(page_count): page = doc[page_num] page_rect = page.rect page_width = page_rect.width page_height = page_rect.height # Calculate text width for spacing estimation text_length = fitz.get_text_length( watermark_text, fontname="helv", fontsize=self.WATERMARK_FONT_SIZE ) # Calculate diagonal for watermark coverage diagonal = math.sqrt(page_width**2 + page_height**2) # Set watermark color with opacity (gray with 30% opacity) color = (0.5, 0.5, 0.5) # Gray # Calculate rotation angle in radians angle_rad = math.radians(self.WATERMARK_ANGLE) # Draw watermark pattern using shape with proper rotation # We use insert_textbox with a morph transform for rotation spacing_x = text_length + self.WATERMARK_SPACING spacing_y = self.WATERMARK_FONT_SIZE + self.WATERMARK_SPACING # Create watermark by drawing rotated text lines # We'll use a simpler approach: draw text and apply rotation via morph shape = page.new_shape() # Calculate grid positions to cover the page when rotated center = fitz.Point(page_width / 2, page_height / 2) # Calculate start and end points for coverage start = -diagonal end = diagonal * 2 y = start row = 0 while y < end: x = start + (spacing_x / 2 if row % 2 else 0) while x < end: # Create text position text_point = fitz.Point(x, y) # Apply rotation around center cos_a = math.cos(angle_rad) sin_a = math.sin(angle_rad) # Translate to origin, rotate, translate back rx = text_point.x - center.x ry = text_point.y - center.y new_x = rx * cos_a - ry * sin_a + center.x new_y = rx * sin_a + ry * cos_a + center.y # Check if the rotated point is within page bounds (with margin) margin = 50 if (-margin <= new_x <= page_width + margin and -margin <= new_y <= page_height + margin): # Insert text using shape with rotation via morph text_rect = fitz.Rect(new_x, new_y, new_x + text_length + 10, new_y + 30) # Use insert_textbox with morph for rotation pivot = fitz.Point(new_x, new_y) morph = (pivot, fitz.Matrix(1, 0, 0, 1, 0, 0).prerotate(self.WATERMARK_ANGLE)) shape.insert_textbox( text_rect, watermark_text, fontname="helv", fontsize=self.WATERMARK_FONT_SIZE, color=color, fill_opacity=self.WATERMARK_OPACITY, morph=morph ) x += spacing_x y += spacing_y row += 1 # Commit the shape drawings shape.commit(overlay=True) # Save to bytes output = io.BytesIO() doc.save(output) doc.close() output.seek(0) logger.info( f"PDF watermark applied successfully for user {user_name} " f"(employee_id: {employee_id}), pages: {page_count}" ) return output.getvalue() def is_supported_image(self, mime_type: str) -> bool: """Check if the mime type is a supported image format.""" supported_types = {'image/png', 'image/jpeg', 'image/jpg'} return mime_type.lower() in supported_types def is_supported_pdf(self, mime_type: str) -> bool: """Check if the mime type is a PDF.""" return mime_type.lower() == 'application/pdf' def supports_watermark(self, mime_type: str) -> bool: """Check if the file type supports watermarking.""" return self.is_supported_image(mime_type) or self.is_supported_pdf(mime_type) # Singleton instance watermark_service = WatermarkService()