"""
SM系列算法加密和文件格式伪装模块
用于将图片、随机种子和bit_position信息加密后，伪装成真正的ZIP或PDF文件
"""
import os
import json
from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT
from gmssl.func import random_hex
import zipfile
import zlib
from io import BytesIO
import base64
from datetime import datetime
import hashlib


class SM4EncryptionWrapper:
    def __init__(self):
        self.crypt_sm4 = CryptSM4()

    def encrypt_data(self, data, key):
        """使用SM4加密数据"""
        # 检查key是否为bytes，如果是则转换为hex字符串
        if isinstance(key, bytes):
            key_str = key.hex()
        else:
            key_str = key
        self.crypt_sm4.set_key(key_str.encode('utf-8'), SM4_ENCRYPT)
        encrypted_data = self.crypt_sm4.crypt_ecb(data.encode('utf-8') if isinstance(data, str) else data)
        return encrypted_data

    def decrypt_data(self, encrypted_data, key):
        """使用SM4解密数据"""
        # 检查key是否为bytes，如果是则转换为hex字符串
        if isinstance(key, bytes):
            key_str = key.hex()
        else:
            key_str = key
        self.crypt_sm4.set_key(key_str.encode('utf-8'), SM4_DECRYPT)
        decrypted_data = self.crypt_sm4.crypt_ecb(encrypted_data)
        return decrypted_data.decode('utf-8')

    def create_real_zip_wrapper(self, image_data, seed, bit_positions, key):
        """创建真正的ZIP格式文件，将加密数据作为ZIP内容"""
        # 组合数据
        combined_data = {
            'image_data': base64.b64encode(image_data).decode('utf-8'),
            'seed': seed,
            'bit_positions': bit_positions
        }
        
        json_data = json.dumps(combined_data)
        
        # 加密数据
        encrypted_data = self.encrypt_data(json_data, key)
        
        # 创建真正的ZIP文件
        zip_buffer = BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
            # 添加加密的有效载荷
            zipf.writestr('payload.bin', encrypted_data)
            # 添加一些看起来正常的文件以增加可信度
            readme_content = "This archive contains important data.\n\n" \
                            "Created with LSB steganography tool.\n" \
                            "Contains hidden information using SM4 encryption."
            zipf.writestr('readme.txt', readme_content)
            info_content = "Version: 1.0\nAuthor: LSB Steganography System\nDate: " + datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            zipf.writestr('info.txt', info_content)
        
        return zip_buffer.getvalue()

    def create_real_pdf_wrapper(self, image_data, seed, bit_positions, key):
        """创建真正的PDF格式文件，将加密数据嵌入其中"""
        # 组合数据
        combined_data = {
            'image_data': base64.b64encode(image_data).decode('utf-8'),
            'seed': seed,
            'bit_positions': bit_positions
        }
        
        json_data = json.dumps(combined_data)
        
        # 加密数据
        encrypted_data = self.encrypt_data(json_data, key)
        
        # 创建一个简单但标准的PDF文档
        pdf_content = "%PDF-1.4\n"
        
        # Object 1: Catalog
        pdf_content += """1 0 obj
<<
  /Type /Catalog
  /Pages 2 0 R
>>
endobj
"""
        
        # Object 2: Pages
        pdf_content += """2 0 obj
<<
  /Type /Pages
  /Count 1
  /Kids [3 0 R]
>>
endobj
"""
        
        # Object 3: Page
        pdf_content += """3 0 obj
<<
  /Type /Page
  /Parent 2 0 R
  /MediaBox [0 0 612 792]
  /Contents 4 0 R
  /Resources <<
    /Font << /F1 5 0 R >>
  >>
>>
endobj
"""
        
        # Object 4: Content Stream - define content_text properly without backslashes in f-string
        content_text = ("Document contains sensitive information.\n\n" +
                       "This file was created with LSB steganography tool.\n" +
                       "Hidden data is securely encrypted using SM4 algorithm.")
        content_safe = content_text.replace('(', '\\(').replace(')', '\\)')
        content_length = len(content_text.encode('latin-1'))
        pdf_content += f"""4 0 obj
<<
  /Length {content_length}
>>
stream
BT
/F1 12 Tf
72 720 TD
({content_safe}) Tj
ET
endstream
endobj
"""
        
        # Object 5: Font
        pdf_content += """5 0 obj
<<
  /Type /Font
  /Subtype /Type1
  /BaseFont /Helvetica
>>
endobj
"""
        
        # Object 6: Our encrypted data as a separate stream
        data_length = len(encrypted_data)
        pdf_content += f"""6 0 obj
<<
  /Type /EmbeddedFile
  /Subtype /application#2Foctet-stream
  /Length {data_length}
>>
stream
"""
        
        # Convert to bytes and append encrypted data
        pdf_bytes = pdf_content.encode('latin-1')
        pdf_bytes += encrypted_data  # Add encrypted data directly as bytes
        pdf_bytes += b"\nendstream\nendobj\n"
        
        # Calculate xref positions
        xref_position = len(pdf_bytes)
        
        # Add xref table and trailer
        xref_table = """xref
0 7
0000000000 65535 f 
0000000017 00000 n 
0000000073 00000 n 
0000000149 00000 n 
0000000274 00000 n 
0000000470 00000 n 
0000000532 00000 n 
trailer
<<
  /Size 7
  /Root 1 0 R
>>
startxref
"""
        
        pdf_bytes += xref_table.encode('latin-1')
        pdf_bytes += str(xref_position).encode('latin-1')
        pdf_bytes += b"\n%%EOF\n"
        
        return pdf_bytes

    def extract_from_real_zip(self, zip_data, key):
        """从真实的ZIP文件中提取并解密数据"""
        zip_buffer = BytesIO(zip_data)
        try:
            with zipfile.ZipFile(zip_buffer, 'r') as zipf:
                # 尝试从不同的文件名中读取加密数据
                try:
                    encrypted_data = zipf.read('payload.bin')
                except KeyError:
                    # 如果没有找到payload.bin，尝试其他可能的文件名
                    encrypted_data = zipf.read('encrypted_payload.bin')
        
            decrypted_json = self.decrypt_data(encrypted_data, key)
            return json.loads(decrypted_json)
        except Exception as e:
            print(f"从ZIP文件提取数据失败: {e}")
            return None

    def extract_from_real_pdf(self, pdf_data, key):
        """从真实的PDF文件中提取并解密数据"""
        try:
            # 尝试解析PDF并找到加密数据
            pdf_str = pdf_data.decode('latin-1', errors='ignore')
            
            # 查找加密数据的位置（在流对象中）
            import re
            
            # 我们的加密数据存储在XObject对象中，尝试匹配这个特定的模式
            # 查找类似 "/Subtype /Image" 的XObject，其中包含了加密数据
            xobject_pattern = r'/Subtype /Image[^>]*/Length \d+\s*>>\s*stream\s*(.*?)\s*endstream'
            match = re.search(xobject_pattern, pdf_str, re.DOTALL)
            
            if match:
                potential_data = match.group(1)
                try:
                    # 尝试将其作为加密数据进行解密
                    decrypted_json = self.decrypt_data(potential_data.encode('latin-1'), key)
                    return json.loads(decrypted_json)
                except Exception as decrypt_error:
                    print(f"PDF解密失败: {decrypt_error}")
            
            # 如果上面的方式不行，尝试通用的流匹配
            stream_matches = re.findall(r'<<[^>]*>>\s*stream\s*(.*?)\s*endstream', pdf_str, re.DOTALL)
            
            # 尝试解密每一个流
            for potential_data in stream_matches:
                try:
                    # 尝试将其作为加密数据进行解密
                    decrypted_json = self.decrypt_data(potential_data.encode('latin-1'), key)
                    return json.loads(decrypted_json)
                except:
                    continue
            
            raise ValueError("未能从PDF中找到加密数据")
        except Exception as e:
            print(f"从PDF文件提取数据失败: {e}")
            return None


def generate_sm4_key_from_password(password):
    """从用户输入的密码生成SM4密钥，支持中文、英文、数字混合"""
    # 使用密码的哈希值来生成固定长度的密钥
    import hashlib
    
    # 将密码转换为UTF-8字节序列，这样可以支持中文字符
    password_bytes = password.encode('utf-8')
    
    # 使用SHA256生成固定长度的哈希值
    hash_obj = hashlib.sha256(password_bytes)
    key = hash_obj.digest()[:16]  # 取前16字节作为SM4密钥
    
    return key


def create_encrypted_wrapper(image_path, output_path, wrapper_format='zip', password=None):
    """创建加密封装文件的便捷函数"""
    if password is None:
        password = input("请输入加密密码: ")
    
    # 读取图像数据
    with open(image_path, 'rb') as f:
        image_data = f.read()
    
    # 读取随机种子和位位置信息
    seed = 0  # 默认值，实际应用中应从相应文件读取
    bit_positions = []  # 默认值，实际应用中应从相应文件读取
    
    # 尝试从文件读取真实值
    if os.path.exists('secret_key.txt'):
        with open('secret_key.txt', 'r') as f:
            seed = int(f.readline().strip())
    
    if os.path.exists('bit_positions.txt'):
        with open('bit_positions.txt', 'r') as f:
            bit_positions = [int(line.strip()) for line in f if line.strip()]
    
    # 创建加密包装器实例
    wrapper = SM4EncryptionWrapper()
    key = generate_sm4_key_from_password(password)
    
    if wrapper_format.lower() == 'zip':
        result = wrapper.create_real_zip_wrapper(image_data, seed, bit_positions, key)
        with open(output_path, 'wb') as f:
            f.write(result)
        print(f"已创建ZIP格式的加密封装文件: {output_path}")
    elif wrapper_format.lower() == 'pdf':
        result = wrapper.create_real_pdf_wrapper(image_data, seed, bit_positions, key)
        with open(output_path, 'wb') as f:
            f.write(result)
        print(f"已创建PDF格式的加密封装文件: {output_path}")
    else:
        raise ValueError("不支持的封装格式，仅支持zip或pdf")


def extract_from_encrypted_wrapper(wrapper_path, wrapper_format='zip', password=None):
    """从加密封装文件中提取数据的便捷函数"""
    if password is None:
        password = input("请输入解密密码: {password}")
    
    key = generate_sm4_key_from_password(password)
    wrapper = SM4EncryptionWrapper()
    
    with open(wrapper_path, 'rb') as f:
        wrapper_data = f.read()
    
    if wrapper_format.lower() == 'zip':
        extracted_data = wrapper.extract_from_real_zip(wrapper_data, key)
    elif wrapper_format.lower() == 'pdf':
        extracted_data = wrapper.extract_from_real_pdf(wrapper_data, key)
    else:
        raise ValueError("不支持的封装格式，仅支持zip或pdf")
    
    if extracted_data is None:
        print("解密失败，请检查密码或文件格式")
        return None
    
    # 将提取的图像数据写入文件
    image_data = base64.b64decode(extracted_data['image_data'])
    with open('extracted_image.bmp', 'wb') as f:
        f.write(image_data)
    
    # 保存提取的种子和位位置
    with open('extracted_secret_key.txt', 'w') as f:
        f.write(str(extracted_data['seed']) + '\n')
    
    with open('extracted_bit_positions.txt', 'w') as f:
        for pos in extracted_data['bit_positions']:
            f.write(str(pos) + '\n')
    
    print(f"已提取图像到 extracted_image.bmp")
    print(f"已提取密钥到 extracted_secret_key.txt")
    print(f"已提取位位置到 extracted_bit_positions.txt")
    
    return extracted_data