app.py

# ShellCode分析工具完整后端实现
# 使用静态文件提供HTML页面，包含登录验证和分析API
import re
import sys
import binascii
import capstone
from collections import defaultdict
from flask import Flask, request, jsonify, session, redirect, url_for, send_from_directory
from openai import OpenAI
import os
import logging
from datetime import datetime, timedelta
from functools import wraps

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建Flask应用
app = Flask(__name__)

# ------------------------------
# 配置
# ------------------------------
app.config['JSON_SORT_KEYS'] = False  # 保持JSON响应字段顺序
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 限制上传大小为16MB
app.config['SECRET_KEY'] = os.environ.get('SECRET_KEY', 'your-secure-secret-key-here')  # 生产环境应使用环境变量
app.config['PERMANENT_SESSION_LIFETIME'] = timedelta(hours=2)  # 会话有效期2小时

# 文件路径配置
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
STATIC_FOLDER = os.path.join(BASE_DIR, 'static')  # 静态文件目录

# 确保静态目录存在
os.makedirs(STATIC_FOLDER, exist_ok=True)

# 固定密码 (在生产环境中应替换为更安全的认证方式)
FIXED_PASSWORD = "Your-Password"  # 请修改此密码


result_dict = {}
def load_system_call_maps(operating_system, arch_mode):
        system_call_dictionary = {
            # Linux系统调用 (x86_64)
            "linux_x64": {
                0x00: "sys_read",
                0x01: "sys_write",
                0x02: "sys_open",
                0x03: "sys_close",
                0x2a: "sys_socket",
                0x21: "sys_connect",
                0x22: "sys_accept",
                0x29: "sys_sendto",
                0x2b: "sys_recvfrom",
                0x3b: "sys_execve",
                0x3c: "sys_exit"
            },
            # Linux系统调用 (x86)
            "linux_x86": {
                0x03: "sys_read",
                0x04: "sys_write",
                0x05: "sys_open",
                0x06: "sys_close",
                0x66: "sys_socketcall",
                0xb0: "sys_execve",
                0xfc: "sys_exit"
            },
            # Windows系统调用 (x64)
            "windows_x64": {
                0x002A: "NtCreateFile",
                0x003A: "NtOpenFile",
                0x0016: "NtReadFile",
                0x0017: "NtWriteFile",
                0x000C: "NtClose",
                0x00AD: "NtCreateThreadEx",
                0x0108: "NtAllocateVirtualMemory",
                0x011C: "NtProtectVirtualMemory"
            },
            # Windows系统调用 (x86)
            "windows_x86": {
                0x002A: "NtCreateFile",
                0x003A: "NtOpenFile",
                0x0016: "NtReadFile",
                0x0017: "NtWriteFile",
                0x000C: "NtClose",
                0x00A2: "NtCreateThread",
                0x0018: "NtAllocateVirtualMemory",
                0x0050: "NtProtectVirtualMemory"
            }
        }
        index_key = "%s_%s" % (str(operating_system), str(arch_mode))
        return system_call_dictionary[index_key]


def disassemble(shellcode, operating_system, index_string, arch_type=capstone.CS_ARCH_X86, arch_mode=capstone.CS_MODE_64, base_addr=0x10000000):
    global result_dict
    statistics = defaultdict(int)
    system_call_maps = load_system_call_maps(operating_system, index_string)
    operator = capstone.Cs(arch_type, arch_mode)
    disassembly = []
    system_calls = []
    try:
        for item in operator.disasm(shellcode, base_addr):
            disassembly.append({
                "address": item.address,
                "mnemonic": item.mnemonic,
                "op_str": item.op_str,
                "bytes": binascii.hexlify(item.bytes).decode()
            })
    except Exception as e:
        raise RuntimeError(f"反汇编失败: {str(e)}")
    statistics["instruction_count"] = len(disassembly)
    control_flow_instructions = {"jmp", "je", "jne", "jz", "jnz", "call", "ret", "cmp", "test"}
    memory_instructions = {"mov", "lea", "push", "pop", "add", "sub", "and", "or", "xor", "inc", "dec"}
    for item in disassembly:
        if item["mnemonic"] in control_flow_instructions:
            statistics["control_flow_instructions"] += 1
        if item["mnemonic"] in memory_instructions:
            statistics["memory_instructions"] += 1
    # 系统调用检测
    system_call_instructions = {"syscall", "int", "int 0x80", "int 0x2e"}
    for i, item in enumerate(disassembly):
        if item["mnemonic"] in system_call_instructions or (item["mnemonic"] == "int" and "0x80" in item["op_str"]):
            # 尝试获取系统调用号
            system_call_num = None
            # 检查前一条指令是否设置系统调用号
            if i > 0 and disassembly[i - 1]["mnemonic"] in ["mov", "xor"]:
                prev_op = disassembly[i - 1]["op_str"]
                if arch_mode == capstone.CS_MODE_64 and "rax" in prev_op:
                    try:
                        system_call_num = int(prev_op.split(",")[-1].strip(), 16)
                    except:
                        pass
                elif arch_mode == capstone.CS_MODE_32 and "eax" in prev_op:
                    try:
                        system_call_num = int(prev_op.split(",")[-1].strip(), 16)
                    except:
                        pass

            call_name = system_call_maps.get(system_call_num, "未知系统调用[%s]" % str(system_call_num))
            system_calls.append({
                "address": item["address"],
                "instruction": f"{item['mnemonic']} {item['op_str']}",
                "syscall_num": system_call_num,
                "name": call_name
            })
            statistics["system_calls"] += 1
    result_dict["instruction_count"] = int(statistics['instruction_count'])
    result_dict["control_flow_instructions"] = int(statistics['control_flow_instructions'])
    result_dict["memory_operations"] = int(statistics['memory_instructions'])
    result_dict["syscall_count"] = int(statistics['syscall_count'])
    result_dict["assembly_code"] = ""
    for item in disassembly:
        result_dict["assembly_code"] += f"0x{item['address']:x}:\t{item['mnemonic']}\t{item['op_str']}\n"


def get_string_from_shellcode(shellcode) -> None:
    global result_dict
    ascii_strings: list = []
    current_str: str = ""
    index_number: int = 1
    ip_pattern = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b")
    url_pattern = re.compile(r"https?://[^\s]+|www\.[^\s]+\.[^\s]+")
    for byte in shellcode:
        if 0x20 <= byte <= 0x7E:
            current_str += chr(byte)
        else:
            if len(current_str) >= 4:
                ascii_strings.append(current_str)
            current_str = ""
    if len(current_str) >= 4:
        ascii_strings.append(current_str)
    result_dict["strings"] = []
    for strings in list(set(ascii_strings)):
        if re.search(ip_pattern, strings):
            result_dict["strings"].append("%s [发现IP地址]" % (str(strings)))
        elif re.search(url_pattern, strings):
            result_dict["strings"].append("%s [发现URL]" % (str(strings)))
        else:
            result_dict["strings"].append(str(strings))
    return None

# ------------------------------
# 辅助函数和装饰器
# ------------------------------

def login_required(f):
    """登录保护装饰器"""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'logged_in' not in session or not session['logged_in']:
            # 对于API请求，返回JSON错误
            if request.path.startswith('/api/'):
                return jsonify({
                    "error": True,
                    "message": "请先登录",
                    "code": "NOT_LOGGED_IN"
                }), 401
            
            # 对于页面请求，重定向到登录页
            return redirect(url_for('login_page'))
        return f(*args, **kwargs)
    return decorated_function

# ------------------------------
# ShellCode分析函数 (用户需要实现)
# ------------------------------
def analyze_shellcode(shellcode, os, architecture):
    """
    ShellCode分析函数，用户需要在此处实现具体的分析逻辑
    
    参数:
        shellcode (str): 输入的ShellCode字符串
        os (str): 操作系统选择 (windows/linux)
        architecture (str): 架构选择 (x86/x64)
    
    返回:
        dict: 包含分析结果的字典
    """
    # 这里仅返回一个示例响应框架，用户需要替换为实际分析逻辑
    logger.info(f"收到分析请求: OS={os}, Architecture={architecture}, ShellCode长度={len(shellcode)}字节")
    
    global  result_dict
    shellcode = binascii.unhexlify(shellcode)
    get_string_from_shellcode(shellcode)
    result_dict["shellcode_length"] = len(shellcode)
    arch_type = capstone.CS_ARCH_X86
    arch_mode = capstone.CS_MODE_32 if architecture == "x86" else capstone.CS_MODE_64
    disassemble(shellcode, os, architecture, arch_type, arch_mode, base_addr=0x10000000)
    
    
    return result_dict

# ------------------------------
# 路由定义
# ------------------------------

@app.route('/shellcodeAnalyzer')
def main_page():
    """主路由，重定向到登录页或分析页"""
    if 'logged_in' in session and session['logged_in']:
        return redirect(url_for('analyzer_page'))
    else:
        return redirect(url_for('login_page'))

@app.route('/shellcodeAnalyzer/login')
def login_page():
    """登录页面"""
    # 如果已登录，重定向到分析页面
    if 'logged_in' in session and session['logged_in']:
        return redirect(url_for('analyzer_page'))
    
    # 返回静态登录页面
    return send_from_directory(STATIC_FOLDER, 'login.html')

@app.route('/shellcodeAnalyzer/analyzer')
@login_required
def analyzer_page():
    """分析工具页面"""
    return send_from_directory(STATIC_FOLDER, 'analyzer.html')

# ------------------------------
# API路由
# ------------------------------

@app.route('/shellcodeAnalyzer/api/login', methods=['POST'])
def api_login():
    """登录API"""
    try:
        data = request.get_json()
        
        if not data or 'password' not in data:
            return jsonify({
                "success": False,
                "message": "缺少密码参数"
            }), 400
        
        # 验证密码
        if data['password'] == FIXED_PASSWORD:
            # 登录成功，设置会话
            session['logged_in'] = True
            session.permanent = True  # 使用永久会话
            
            logger.info("用户登录成功")
            return jsonify({
                "success": True,
                "message": "登录成功"
            })
        else:
            logger.warning("用户登录失败：密码错误")
            return jsonify({
                "success": False,
                "message": "密码错误，请重试"
            }), 401
            
    except Exception as e:
        logger.error(f"登录过程出错: {str(e)}")
        return jsonify({
            "success": False,
            "message": "登录过程出错，请稍后重试"
        }), 500

@app.route('/shellcodeAnalyzer/api/logout', methods=['POST'])
def api_logout():
    """登出API"""
    # 清除会话
    session.pop('logged_in', None)
    logger.info("用户已登出")
    return jsonify({
        "success": True,
        "message": "登出成功"
    })

@app.route('/shellcodeAnalyzer/api/check-login', methods=['GET'])
def api_check_login():
    """检查登录状态API"""
    return jsonify({
        "loggedIn": 'logged_in' in session and session['logged_in']
    })




@app.route('/shellcodeAnalyzer/api/analyze', methods=['POST'])
@login_required
def api_analyze():
    """ShellCode分析API接口"""
    start_time = datetime.now()
    global result_dict
    try:
        # 获取请求数据
        data = request.get_json()
        
        # 验证请求数据
        if not data or 'shellcode' not in data:
            return jsonify({
                "error": True,
                "message": "请求缺少必要参数: shellcode",
                "code": "MISSING_PARAMETER"
            }), 400
        
        # 提取参数，设置默认值
        shellcode = data.get('shellcode', '')
        os = data.get('os', 'windows')
        architecture = data.get('architecture', 'x64')
        result_dict["os_version"] = os
        result_dict["architecture"] = architecture
        # 调用分析函数（用户需要实现此函数）
        result = analyze_shellcode(shellcode, os, architecture)
        
        # 记录分析耗时
        elapsed_time = (datetime.now() - start_time).total_seconds()
        logger.info(f"分析完成，耗时: {elapsed_time:.4f}秒")
        
        # 返回分析结果
        return jsonify(result_dict)
        
    except Exception as e:
        # 记录异常信息
        logger.error(f"分析过程发生错误: {str(e)}", exc_info=True)
        
        # 返回错误响应
        return jsonify({
            "error": True,
            "message": f"分析过程发生错误: {str(e)}",
            "code": "ANALYSIS_ERROR"
        }), 500



def get_ai_analysis(shellcode, os, architecture):
    client = OpenAI(api_key="your_api_token", base_url="https://api.siliconflow.cn/v1")
    response = client.chat.completions.create(
        model="your_model_name",
        messages=[{
            "role": "user",
            "content": "你是一名系统安全专家，请分析这段%s %s 的shellcode 并给出解读:  %s 纯文字形式给出你的分析依据和结论，不要有任何markdown或者格式，纯文本输出即可。" % (str(os), str(architecture), str(shellcode))
        }],
        temperature=0.7,
        max_tokens=4096
    )
    ai_result_string = response.to_dict()["choices"][0]["message"]["content"]
    return {"ai_result": ai_result_string}


@app.route('/shellcodeAnalyzer/api/ai-analysis', methods=['POST'])
@login_required
def api_ai_analysis():
    """AI分析API接口 - 可独立调用"""
    try:
        # 获取请求数据
        data = request.get_json()
        
        # 验证请求数据
        if not data or 'shellcode' not in data:
            return jsonify({
                "error": True,
                "message": "请求缺少必要参数: shellcode",
                "code": "MISSING_PARAMETER"
            }), 400
        
        # 提取参数，设置默认值
        shellcode = data.get('shellcode', '')
        os = data.get('os', 'windows')
        architecture = data.get('architecture', 'x64')
        
        
        # 获取AI分析结果（可独立调用）
        ai_result = get_ai_analysis(shellcode, os, architecture)
        
        # 返回AI分析结果
        return jsonify(ai_result)
        
    except Exception as e:
        # 记录异常信息
        logger.error(f"AI分析过程发生错误: {str(e)}", exc_info=True)
        
        # 返回错误响应
        return jsonify({
            "error": True,
            "message": f"AI分析过程发生错误: {str(e)}",
            "code": "AI_ANALYSIS_ERROR"
        }), 500

# ------------------------------
# 静态文件路由 (用于访问CSS/JS等资源)
# ------------------------------
@app.route('/shellcodeAnalyzer/static/<path:filename>')
def static_files(filename):
    """提供静态资源访问"""
    return send_from_directory(STATIC_FOLDER, filename)

# ------------------------------
# 错误处理
# ------------------------------

@app.errorhandler(401)
def unauthorized(error):
    """未授权访问错误处理"""
    return jsonify({
        "error": True,
        "message": "请先登录",
        "code": "NOT_LOGGED_IN"
    }), 401

@app.errorhandler(404)
def not_found(error):
    """404错误处理"""
    return jsonify({
        "error": True,
        "message": "请求的资源不存在",
        "code": "RESOURCE_NOT_FOUND"
    }), 404

@app.errorhandler(413)
def request_entity_too_large(error):
    """请求体过大错误处理"""
    return jsonify({
        "error": True,
        "message": "请求体过大，超过16MB限制",
        "code": "REQUEST_TOO_LARGE"
    }), 413

# ------------------------------
# 应用入口
# ------------------------------
if __name__ == '__main__':
    # 显示项目结构提示
    logger.info(f"项目结构:")
    logger.info(f"当前文件: {os.path.abspath(__file__)}")
    logger.info(f"静态文件目录: {STATIC_FOLDER}")
    logger.info(f"请将HTML文件放入静态目录，结构如下:")
    logger.info(f"static/")
    logger.info(f"├── login.html      - 登录页面")
    logger.info(f"└── analyzer.html   - 分析工具页面")
    logger.info(f"默认登录密码: {FIXED_PASSWORD} (请在生产环境中修改)")
    
    # 启动Flask应用
    app.run(
        host='0.0.0.0',  # 允许外部访问
        port= 5008,        # 默认端口
        debug=False        # 开发模式，生产环境应设为False
    )