[ { "path": ".github/FUNDING.yml", "content": "# These are supported funding model platforms\n\ngithub: # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]\npatreon: #\nopen_collective: # Replace with a single Open Collective username\nko_fi: # Replace with a single Ko-fi username\ntidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel\ncommunity_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry\nliberapay: # Replace with a single Liberapay username\nissuehunt: # Replace with a single IssueHunt username\notechie: # Replace with a single Otechie username\ncustom: ['http://pic.giantbranch.cn/pic/1551450728861.jpg']\n" }, { "path": "README.md", "content": "# IDAPython mipsAudit\n\n## 简介\n\n这是一个简单的IDAPython脚本。\n\n进一步来说是MIPS静态汇编审计辅助脚本。\n\n可能会有bug，欢迎大家完善。\n\n> **v3.0 更新**: 支持 IDA 7.x，8.x API 和 Python 3，新增多项高级漏洞检测功能。\n\n## 功能\n\n### 基础功能\n\n1. 找到危险函数的调用处，并且高亮该行（也可以下断点,这个需要自己去源码看吧）\n2. 给参数赋值处加上注释\n3. 最后以表格的形式输出函数名，调用地址，参数，还有当前函数的缓冲区大小\n\n**大家双击addr那一列的地址，即可跳到对应的地址处**\n\n![17cc62c98820974f8c759dc086dd5acb](17cc62c98820974f8c759dc086dd5acb.png)\n\n![28069d48cf3f357dd83e42406e10d980](28069d48cf3f357dd83e42406e10d980.png)\n\n### v3.0 新增功能\n\n#### 高级漏洞检测\n\n| 检测类型 | 说明 |\n|---------|------|\n| **命令注入检测** | 追踪 system/popen/execve 参数是否来自外部输入 |\n| **栈溢出检测** | 比较目标缓冲区大小与源数据长度 |\n| **格式化字符串漏洞** | 检测 %n 写入原语和用户可控格式字符串 |\n| **整数溢出检测** | 检测 malloc/calloc 的 size 参数来源和算术运算 |\n| **双重释放检测** | 追踪 free() 调用，检测同一指针多次释放 |\n| **数据流分析** | 追踪 read/recv 返回数据流向危险函数 |\n| **Wrapper函数识别** | 识别封装了危险函数的自定义函数 |\n\n#### 风险等级高亮\n\n| 颜色 | 等级 | 说明 |\n|------|------|------|\n| 红色 | HIGH | 需要立即关注 |\n| 橙色 | MEDIUM | 需要人工复核 |\n| 绿色 | LOW | 信息提示 |\n\n#### 结果导出\n\n自动生成带时间戳的 HTML 报告，保存到 IDB 文件所在目录：\n```\nmipsAudit_results_20260129_143052.html\n```\n\n#### 配置文件支持\n\n支持通过 `mipsAudit_config.json` 扩展自定义函数列表：\n```json\n{\n \"dangerous_functions\": [\"custom_strcpy\"],\n \"command_execution_function\": [\"custom_exec\"],\n \"_comment\": \"扩展默认函数列表\"\n}\n```\n\n## 审计的危险函数\n\n```python\ndangerous_functions = [\n \"strcpy\", \"strcat\", \"sprintf\", \"read\", \"getenv\",\n \"gets\", \"scanf\", \"vscanf\", \"realpath\", \"access\", \"stat\", \"lstat\"\n]\n\nattention_function = [\n \"memcpy\", \"strncpy\", \"sscanf\", \"strncat\", \"snprintf\",\n \"vprintf\", \"printf\", \"fprintf\", \"vfprintf\", \"vsprintf\",\n \"vsnprintf\", \"syslog\", \"memmove\", \"bcopy\"\n]\n\ncommand_execution_function = [\n \"system\", \"execve\", \"popen\", \"unlink\",\n \"execl\", \"execle\", \"execlp\", \"execv\", \"execvp\",\n \"dlopen\", \"mmap\", \"mprotect\"\n]\n\nmemory_alloc_functions = [\n \"malloc\", \"calloc\", \"realloc\", \"memalign\",\n \"valloc\", \"pvalloc\", \"aligned_alloc\", \"mmap\"\n]\n\nmemory_free_functions = [\n \"free\", \"cfree\", \"munmap\"\n]\n```\n\n## 运行流程\n\n```\nPHASE 1: Basic Function Audit # 基础危险函数审计\nPHASE 2: Enhanced Vulnerability Detection # 增强漏洞检测\nPHASE 3: Advanced Analysis # 高级分析（数据流、Wrapper识别）\nPHASE 4: Results Summary & Export # 结果汇总与导出\n```\n\n## 使用\n\n### 环境要求\n\n- IDA Pro 7.0+\n- Python 3.x（IDA 内置）\n\n### 运行方式\n\nFile - Script file\n\n![1561006651468](./1561006651468.png)\n\n选择mipsAudit.py\n\n![1561006737134](./1561006737134.png)\n\n即可看到效果\n\n![mipsAudit](./mipsAudit.png)\n\n双击地址即可跳到对应的代码处\n\n![1561006887117](./1561006887117.png)\n\n## 更新日志\n\n### v3.0 (2026-01)\n- 支持 IDA 7.x+ API\n- 迁移至 Python 3\n- 新增格式化字符串漏洞检测（%n 检测）\n- 新增命令注入参数来源追踪\n- 新增栈溢出缓冲区大小比较\n- 新增整数溢出检测（malloc/calloc size 参数）\n- 新增双重释放检测\n- 新增数据流分析（read/recv 返回值追踪）\n- 新增 Wrapper 函数识别\n- 新增基本块分析（修复跨基本块误判）\n- 新增 HTML 报告导出（带时间戳）\n- 新增外部 JSON 配置文件支持\n- 新增扫描进度显示\n- 扩展危险函数列表\n\n### v1.0 (2018-05)\n- 初始版本 by giantbranch\n" }, { "path": "mipsAudit.py", "content": "# -*- coding: utf-8 -*-\n\n# reference\n# 《ida pro 权威指南》\n# 《python 灰帽子》\n# 《家用路由器0day漏洞挖掘》\n# https://github.com/wangzery/SearchOverflow/blob/master/SearchOverflow.py\n\n# Updated for IDA 7.x+ API and Python 3\n# Enhanced with advanced vulnerability detection v3.0\n\nimport idc\nimport idaapi\nimport idautils\nfrom prettytable import PrettyTable\nfrom collections import defaultdict\nimport json\nimport csv\nimport os\nimport re\nfrom datetime import datetime\n\n# IDA 7.x+ uses BADADDR from idaapi\nBADADDR = idaapi.BADADDR\n\n\nDEBUG = True\n\n# ============================================================\n# Configuration - Can be overridden by external config file\n# ============================================================\n\nCONFIG_FILE = \"mipsAudit_config.json\"\n\n# Default function lists (can be extended via config file)\ndangerous_functions = [\n \"strcpy\", \n \"strcat\", \n \"sprintf\",\n \"read\", \n \"getenv\",\n \"gets\", # No boundary check - extremely dangerous\n \"scanf\", # Format input vulnerability\n \"vscanf\",\n \"realpath\", # Path traversal\n \"access\", # TOCTOU race condition\n \"stat\", # TOCTOU race condition\n \"lstat\",\n]\n\nattention_function = [\n \"memcpy\",\n \"strncpy\",\n \"sscanf\", \n \"strncat\", \n \"snprintf\",\n \"vprintf\", \n \"printf\",\n \"fprintf\",\n \"vfprintf\",\n \"vsprintf\",\n \"vsnprintf\",\n \"syslog\",\n \"memmove\",\n \"bcopy\",\n]\n\ncommand_execution_function = [\n \"system\", \n \"execve\",\n \"popen\",\n \"unlink\",\n \"execl\",\n \"execle\", \n \"execlp\",\n \"execv\",\n \"execvp\",\n \"dlopen\",\n \"mmap\", # Memory mapping\n \"mprotect\", # Change memory protection\n]\n\n# External input source functions (for taint tracking)\nexternal_input_functions = [\n \"getenv\",\n \"read\",\n \"recv\",\n \"recvfrom\",\n \"recvmsg\",\n \"fgets\",\n \"fread\",\n \"fgetc\",\n \"gets\",\n \"scanf\",\n \"fscanf\",\n \"getchar\",\n \"getc\",\n \"fgetws\",\n \"getwchar\",\n \"getline\",\n \"getdelim\",\n \"socket\",\n \"accept\",\n \"gethostbyname\",\n]\n\n# Memory management functions\nmemory_alloc_functions = [\n \"malloc\",\n \"calloc\",\n \"realloc\",\n \"memalign\",\n \"valloc\",\n \"pvalloc\",\n \"aligned_alloc\",\n \"mmap\",\n]\n\nmemory_free_functions = [\n \"free\",\n \"cfree\",\n \"munmap\",\n]\n\n# Format string functions (for %n detection)\nformat_string_functions = [\n \"printf\",\n \"fprintf\", \n \"sprintf\",\n \"snprintf\",\n \"vprintf\",\n \"vfprintf\",\n \"vsprintf\",\n \"vsnprintf\",\n \"syslog\",\n]\n\n# describe arg num of function\none_arg_function = [\n \"getenv\",\n \"system\",\n \"unlink\",\n \"free\",\n \"cfree\",\n \"malloc\",\n \"gets\",\n]\n\ntwo_arg_function = [\n \"strcpy\", \n \"strcat\",\n \"popen\",\n \"calloc\",\n \"dlopen\",\n \"fgets\",\n \"access\",\n \"stat\",\n \"lstat\",\n \"realpath\",\n \"mprotect\",\n]\n\nthree_arg_function = [\n \"strncpy\",\n \"strncat\", \n \"memcpy\",\n \"memmove\",\n \"bcopy\",\n \"execve\",\n \"read\",\n \"recv\",\n \"fread\",\n \"realloc\",\n]\n\nformat_function_offset_dict = {\n \"sprintf\": 1,\n \"sscanf\": 1,\n \"snprintf\": 2,\n \"vprintf\": 0,\n \"printf\": 0,\n \"fprintf\": 1,\n \"vfprintf\": 1,\n \"vsprintf\": 1,\n \"vsnprintf\": 2,\n \"syslog\": 1,\n \"scanf\": 0,\n}\n\n# Risk level colors\nRISK_HIGH = 0x0000ff # Red\nRISK_MEDIUM = 0x00a5ff # Orange\nRISK_LOW = 0x00ff00 # Green\nRISK_INFO = 0xffff00 # Cyan\n\n# ============================================================\n# Enhanced Analysis - Data Structures\n# ============================================================\n\n# Store function call information for cross-reference analysis\ntaint_sources = {} # addr -> function_name (external input sources)\nfree_calls = defaultdict(list) # func_addr -> [(call_addr, arg_info), ...]\naudit_results = [] # Store all findings for export\nwrapper_functions = {} # Detected wrapper functions\ndata_flow_graph = defaultdict(list) # addr -> [(dest_addr, dest_func), ...]\n\n# Progress tracking\ntotal_functions = 0\nprocessed_functions = 0\n\n\n# ============================================================\n# Configuration File Support\n# ============================================================\n\ndef get_output_dir():\n \"\"\"Get output directory - uses IDB directory in IDA environment\"\"\"\n try:\n idb_path = idc.get_idb_path()\n if idb_path:\n return os.path.dirname(idb_path)\n except:\n pass\n return os.getcwd()\n\n\ndef load_config():\n \"\"\"Load configuration from external JSON file\"\"\"\n global dangerous_functions, attention_function, command_execution_function\n global external_input_functions, memory_alloc_functions, memory_free_functions\n global format_string_functions\n \n config_path = os.path.join(get_output_dir(), CONFIG_FILE)\n if os.path.exists(config_path):\n try:\n with open(config_path, 'r', encoding='utf-8') as f:\n config = json.load(f)\n \n # Extend function lists from config\n if 'dangerous_functions' in config:\n dangerous_functions.extend(config['dangerous_functions'])\n if 'attention_function' in config:\n attention_function.extend(config['attention_function'])\n if 'command_execution_function' in config:\n command_execution_function.extend(config['command_execution_function'])\n if 'external_input_functions' in config:\n external_input_functions.extend(config['external_input_functions'])\n if 'memory_alloc_functions' in config:\n memory_alloc_functions.extend(config['memory_alloc_functions'])\n if 'memory_free_functions' in config:\n memory_free_functions.extend(config['memory_free_functions'])\n if 'format_string_functions' in config:\n format_string_functions.extend(config['format_string_functions'])\n \n print(\"[*] Loaded configuration from %s\" % config_path)\n return True\n except Exception as e:\n print(\"[!] Error loading config: %s\" % str(e))\n return False\n\n\ndef save_default_config():\n \"\"\"Save current configuration as default config file\"\"\"\n config_path = os.path.join(get_output_dir(), CONFIG_FILE)\n config = {\n \"dangerous_functions\": [],\n \"attention_function\": [],\n \"command_execution_function\": [],\n \"external_input_functions\": [],\n \"memory_alloc_functions\": [],\n \"memory_free_functions\": [],\n \"format_string_functions\": [],\n \"_comment\": \"Add custom functions to extend the default lists\"\n }\n try:\n with open(config_path, 'w', encoding='utf-8') as f:\n json.dump(config, f, indent=4)\n print(\"[*] Created default config file: %s\" % config_path)\n except Exception as e:\n print(\"[!] Error saving config: %s\" % str(e))\n\n\n# ============================================================\n# Progress Display\n# ============================================================\n\ndef show_progress(current, total, prefix=\"Progress\"):\n \"\"\"Display progress bar in IDA output\"\"\"\n if total == 0:\n return\n percent = (current * 100) // total\n bar_len = 30\n filled = (current * bar_len) // total\n bar = '=' * filled + '-' * (bar_len - filled)\n print(\"\\r%s: [%s] %d%% (%d/%d)\" % (prefix, bar, percent, current, total), end='')\n if current == total:\n print() # New line when complete\n\n\n# ============================================================\n# Helper Functions\n# ============================================================\n\ndef printFunc(func_name):\n string1 = \"========================================\"\n string2 = \"========== Auditing \" + func_name + \" \"\n strlen = len(string1) - len(string2)\n return string1 + \"\\n\" + string2 + '=' * strlen + \"\\n\" + string1\n\ndef getFuncAddr(func_name):\n func_addr = idc.get_name_ea_simple(func_name)\n if func_addr != BADADDR:\n print(printFunc(func_name))\n return func_addr\n return False\n\ndef getFormatString(addr):\n op_num = 1\n # idc.get_operand_type Return value\n # o_void 0 // No Operand\n # o_reg 1 // General Register (al, ax, es, ds...) reg\n # o_mem 2 // Direct Memory Reference (DATA) addr\n # o_phrase 3 // Memory Ref [Base Reg + Index Reg] phrase\n # o_displ 4 // Memory Reg [Base Reg + Index Reg + Displacement] phrase+addr\n # o_imm 5 // Immediate Value value\n # o_far 6 // Immediate Far Address (CODE) addr\n # o_near 7 // Immediate Near Address (CODE) addr\n # o_idpspec0 8 // IDP specific type\n # o_idpspec1 9 // IDP specific type\n # o_idpspec2 10 // IDP specific type\n # o_idpspec3 11 // IDP specific type\n # o_idpspec4 12 // IDP specific type\n # o_idpspec5 13 // IDP specific type\n if idc.get_operand_type(addr, op_num) != 5:\n op_num = op_num + 1\n if idc.get_operand_type(addr, op_num) != 5:\n return \"get fail\"\n op_string = idc.print_operand(addr, op_num).split(\" \")[0].split(\"+\")[0].split(\"-\")[0].replace(\"(\", \"\")\n string_addr = idc.get_name_ea_simple(op_string)\n if string_addr == BADADDR:\n return \"get fail\"\n string_content = idc.get_strlit_contents(string_addr)\n if string_content is None:\n return \"get fail\"\n if isinstance(string_content, bytes):\n string_content = string_content.decode('utf-8', errors='replace')\n return [string_addr, string_content]\n\n\ndef getArgAddr(start_addr, regNum):\n mipscondition = [\"bn\", \"be\" , \"bg\", \"bl\"]\n scan_deep = 50\n count = 0\n reg = \"$a\" + str(regNum)\n # try to get in the next (code references from this address)\n for next_addr in idautils.CodeRefsFrom(start_addr, 0):\n if next_addr != BADADDR and reg == idc.print_operand(next_addr, 0):\n return next_addr\n # try to get before (code references to this address)\n before_addr = start_addr\n for ref_addr in idautils.CodeRefsTo(start_addr, 0):\n before_addr = ref_addr\n break\n if before_addr == start_addr:\n before_addr = idc.prev_head(start_addr)\n \n while before_addr != BADADDR:\n if reg == idc.print_operand(before_addr, 0):\n Mnemonics = idc.print_insn_mnem(before_addr)\n if Mnemonics[0:2] in mipscondition:\n pass\n elif Mnemonics[0:1] == \"j\":\n pass\n else:\n return before_addr\n count = count + 1\n if count > scan_deep:\n break \n before_addr = idc.prev_head(before_addr)\n return BADADDR\n\n\ndef getArg(start_addr, regNum):\n mipsmov = [\"move\", \"lw\", \"li\", \"lb\", \"lui\", \"lhu\", \"lbu\", \"la\"]\n arg_addr = getArgAddr(start_addr, regNum)\n if arg_addr != BADADDR:\n Mnemonics = idc.print_insn_mnem(arg_addr) \n if Mnemonics[0:3] == \"add\":\n if idc.print_operand(arg_addr, 2) == \"\":\n arg = idc.print_operand(arg_addr, 0) + \"+\" + idc.print_operand(arg_addr, 1)\n else:\n arg = idc.print_operand(arg_addr, 1) + \"+\" + idc.print_operand(arg_addr, 2)\n elif Mnemonics[0:3] == \"sub\":\n if idc.print_operand(arg_addr, 2) == \"\":\n arg = idc.print_operand(arg_addr, 0) + \"-\" + idc.print_operand(arg_addr, 1)\n else:\n arg = idc.print_operand(arg_addr, 1) + \"-\" + idc.print_operand(arg_addr, 2)\n elif Mnemonics in mipsmov:\n arg = idc.print_operand(arg_addr, 1) \n else:\n arg = idc.GetDisasm(arg_addr).split(\"#\")[0]\n idc.set_cmt(arg_addr, \"addr: 0x%x \" % start_addr + \"-------> arg\" + str((int(regNum)+1)) + \" : \" + arg, 0)\n return arg\n else:\n return \"get fail\"\n\ndef audit(func_name):\n func_addr = getFuncAddr(func_name) \n if func_addr == False:\n return False\n\n # get arg num and set table\n if func_name in one_arg_function:\n arg_num = 1\n elif func_name in two_arg_function:\n arg_num = 2\n elif func_name in three_arg_function:\n arg_num = 3\n elif func_name in format_function_offset_dict:\n arg_num = format_function_offset_dict[func_name] + 1\n else:\n print(\"The %s function didn't write in the describe arg num of function array, please add it to, such as add to `two_arg_function` array\" % func_name)\n return\n # mispcall = [\"jal\", \"jalr\", \"bal\", \"jr\"]\n table_head = [\"func_name\", \"addr\"]\n for num in range(0, arg_num):\n table_head.append(\"arg\"+str(num+1))\n if func_name in format_function_offset_dict:\n table_head.append(\"format&value[string_addr, num of '%', fmt_arg...]\")\n table_head.append(\"local_buf_size\")\n table = PrettyTable(table_head)\n\n # get references to function (xrefs)\n for call_addr in idautils.CodeRefsTo(func_addr, 0):\n # set color - green (red=0x0000ff, blue=0xff0000)\n idc.set_color(call_addr, idc.CIC_ITEM, 0x00ff00)\n # set break point\n # idc.add_bpt(call_addr)\n # idc.del_bpt(call_addr)\n\n Mnemonics = idc.print_insn_mnem(call_addr)\n if Mnemonics[0:1] == \"j\" or Mnemonics[0:1] == \"b\":\n if func_name in format_function_offset_dict:\n info = auditFormat(call_addr, func_name, arg_num)\n else:\n info = auditAddr(call_addr, func_name, arg_num)\n table.add_row(info)\n print(table)\n # data_addr = DfirstB(func_addr)\n # while data_addr != BADADDR:\n # Mnemonics = GetMnem(data_addr)\n # if DEBUG:\n # print \"Data Mnemonics : %s\" % GetMnem(data_addr)\n # print \"Data addr : 0x %s\" % data_addr\n # data_addr = DnextB(func_addr, data_addr)\n\ndef auditAddr(call_addr, func_name, arg_num):\n addr = \"0x%x\" % call_addr\n ret_list = [func_name, addr]\n # local buf size\n local_buf_size = idc.get_func_attr(call_addr, idc.FUNCATTR_FRSIZE)\n if local_buf_size == BADADDR:\n local_buf_size = \"get fail\"\n else:\n local_buf_size = \"0x%x\" % local_buf_size\n # get arg\n for num in range(0, arg_num):\n ret_list.append(getArg(call_addr, num)) \n ret_list.append(local_buf_size)\n return ret_list\n\ndef auditFormat(call_addr, func_name, arg_num):\n addr = \"0x%x\" % call_addr\n ret_list = [func_name, addr]\n # local buf size\n local_buf_size = idc.get_func_attr(call_addr, idc.FUNCATTR_FRSIZE)\n if local_buf_size == BADADDR:\n local_buf_size = \"get fail\"\n else:\n local_buf_size = \"0x%x\" % local_buf_size\n # get arg\n for num in range(0, arg_num):\n ret_list.append(getArg(call_addr, num)) \n arg_addr = getArgAddr(call_addr, format_function_offset_dict[func_name])\n string_and_addr = getFormatString(arg_addr)\n format_and_value = []\n if string_and_addr == \"get fail\":\n ret_list.append(\"get fail\")\n else:\n string_addr = \"0x%x\" % string_and_addr[0]\n format_and_value.append(string_addr)\n string = string_and_addr[1]\n fmt_num = string.count(\"%\")\n format_and_value.append(fmt_num)\n # mips arg reg is from a0 to a3\n if fmt_num > 3:\n fmt_num = fmt_num - format_function_offset_dict[func_name] - 1\n for num in range(0, fmt_num):\n if arg_num + num > 3:\n break\n format_and_value.append(getArg(call_addr, arg_num + num))\n ret_list.append(format_and_value)\n ret_list.append(local_buf_size)\n return ret_list\n\n# ============================================================\n# Enhanced Detection Functions\n# ============================================================\n\ndef getCallingFunction(addr):\n \"\"\"Get the function containing the given address\"\"\"\n func = idaapi.get_func(addr)\n if func:\n return func.start_ea\n return BADADDR\n\ndef traceArgSource(start_addr, regNum, depth=10):\n \"\"\"\n Trace the source of an argument to detect external input\n Returns: (source_type, source_info)\n source_type: 'external_input', 'static_string', 'stack_var', 'unknown'\n \"\"\"\n if depth <= 0:\n return ('unknown', None)\n \n mipsmov = [\"move\", \"lw\", \"li\", \"lb\", \"lui\", \"lhu\", \"lbu\", \"la\"]\n arg_addr = getArgAddr(start_addr, regNum)\n \n if arg_addr == BADADDR:\n return ('unknown', None)\n \n Mnemonics = idc.print_insn_mnem(arg_addr)\n operand1 = idc.print_operand(arg_addr, 1)\n \n # Check if loading from a static string\n if Mnemonics in [\"la\", \"lui\", \"li\"]:\n str_addr = idc.get_name_ea_simple(operand1.split(\"+\")[0].split(\"-\")[0].replace(\"(\", \"\"))\n if str_addr != BADADDR:\n str_content = idc.get_strlit_contents(str_addr)\n if str_content:\n if isinstance(str_content, bytes):\n str_content = str_content.decode('utf-8', errors='replace')\n return ('static_string', {'addr': str_addr, 'value': str_content, 'len': len(str_content)})\n \n # Check if it's a stack variable\n if \"sp\" in operand1 or \"fp\" in operand1:\n return ('stack_var', {'operand': operand1})\n \n # Check if moved from another register - trace further\n if Mnemonics == \"move\":\n src_reg = operand1\n if src_reg.startswith(\"$v\"): # Return value from function call\n # Look for preceding function call\n scan_addr = idc.prev_head(arg_addr)\n scan_count = 0\n while scan_addr != BADADDR and scan_count < 20:\n mnem = idc.print_insn_mnem(scan_addr)\n if mnem in [\"jal\", \"jalr\"]:\n # Get the called function name\n call_target = idc.print_operand(scan_addr, 0)\n if call_target in external_input_functions:\n return ('external_input', {'function': call_target, 'addr': scan_addr})\n break\n scan_count += 1\n scan_addr = idc.prev_head(scan_addr)\n \n return ('unknown', None)\n\n\ndef checkCommandInjection(call_addr, func_name):\n \"\"\"\n Check if command execution function has external input as argument\n Returns risk assessment\n \"\"\"\n results = []\n \n # For system(), popen() - check first argument (command string)\n if func_name in [\"system\"]:\n source_type, source_info = traceArgSource(call_addr, 0, depth=15)\n if source_type == 'external_input':\n results.append({\n 'risk': 'HIGH',\n 'issue': 'Command injection - arg from %s()' % source_info['function'],\n 'detail': source_info\n })\n elif source_type == 'stack_var':\n results.append({\n 'risk': 'MEDIUM', \n 'issue': 'Command from stack variable (needs manual review)',\n 'detail': source_info\n })\n elif source_type == 'static_string':\n results.append({\n 'risk': 'LOW',\n 'issue': 'Static command string',\n 'detail': source_info\n })\n \n elif func_name in [\"popen\"]:\n source_type, source_info = traceArgSource(call_addr, 0, depth=15)\n if source_type == 'external_input':\n results.append({\n 'risk': 'HIGH',\n 'issue': 'Popen injection - arg from %s()' % source_info['function'],\n 'detail': source_info\n })\n \n elif func_name in [\"execve\", \"execl\", \"execle\", \"execlp\", \"execv\", \"execvp\"]:\n source_type, source_info = traceArgSource(call_addr, 0, depth=15)\n if source_type == 'external_input':\n results.append({\n 'risk': 'HIGH',\n 'issue': 'Exec injection - arg from %s()' % source_info['function'],\n 'detail': source_info\n })\n \n return results\n\n\ndef checkStackOverflow(call_addr, func_name):\n \"\"\"\n Check for potential stack buffer overflow\n Compares destination buffer size with source length\n \"\"\"\n results = []\n local_buf_size = idc.get_func_attr(call_addr, idc.FUNCATTR_FRSIZE)\n \n if func_name in [\"strcpy\", \"strcat\"]:\n # Check source (arg1) for static string length\n source_type, source_info = traceArgSource(call_addr, 1, depth=15)\n if source_type == 'static_string' and source_info:\n src_len = source_info.get('len', 0)\n if local_buf_size != BADADDR and src_len > 0:\n if src_len > local_buf_size:\n results.append({\n 'risk': 'HIGH',\n 'issue': 'Buffer overflow: src_len(%d) > frame_size(0x%x)' % (src_len, local_buf_size),\n 'detail': source_info\n })\n elif src_len > local_buf_size // 2:\n results.append({\n 'risk': 'MEDIUM',\n 'issue': 'Potential overflow: src_len(%d) close to frame_size(0x%x)' % (src_len, local_buf_size),\n 'detail': source_info\n })\n elif source_type == 'external_input':\n results.append({\n 'risk': 'HIGH',\n 'issue': '%s with external input from %s()' % (func_name, source_info['function']),\n 'detail': source_info\n })\n \n elif func_name in [\"sprintf\"]:\n # Check format string for potential overflow\n arg_addr = getArgAddr(call_addr, 1) # format string is arg1\n fmt_info = getFormatString(arg_addr)\n if fmt_info != \"get fail\":\n fmt_str = fmt_info[1]\n # Check for %s without width limit\n if '%s' in fmt_str:\n results.append({\n 'risk': 'MEDIUM',\n 'issue': 'sprintf with %%s (unbounded string copy)',\n 'detail': {'format': fmt_str}\n })\n \n elif func_name in [\"memcpy\", \"strncpy\", \"strncat\"]:\n # Check size parameter (arg2)\n source_type, source_info = traceArgSource(call_addr, 2, depth=15)\n if source_type == 'external_input':\n results.append({\n 'risk': 'HIGH',\n 'issue': '%s size from external input %s()' % (func_name, source_info['function']),\n 'detail': source_info\n })\n \n return results\n\n\ndef checkIntegerOverflow(call_addr, func_name):\n \"\"\"\n Check for potential integer overflow in memory allocation\n \"\"\"\n results = []\n \n if func_name in [\"malloc\", \"calloc\", \"realloc\", \"memalign\"]:\n # Check size argument\n size_arg_idx = 0 if func_name == \"malloc\" else 1\n if func_name == \"calloc\":\n # calloc(count, size) - check both arguments\n for idx in [0, 1]:\n source_type, source_info = traceArgSource(call_addr, idx, depth=15)\n if source_type == 'external_input':\n results.append({\n 'risk': 'HIGH',\n 'issue': 'calloc arg%d from external input %s()' % (idx, source_info['function']),\n 'detail': source_info\n })\n else:\n source_type, source_info = traceArgSource(call_addr, size_arg_idx, depth=15)\n if source_type == 'external_input':\n results.append({\n 'risk': 'HIGH',\n 'issue': '%s size from external input %s()' % (func_name, source_info['function']),\n 'detail': source_info\n })\n \n # Check for arithmetic operations on size (potential integer overflow)\n arg_addr = getArgAddr(call_addr, size_arg_idx)\n if arg_addr != BADADDR:\n # Scan backward for arithmetic operations\n scan_addr = arg_addr\n scan_count = 0\n while scan_addr != BADADDR and scan_count < 10:\n mnem = idc.print_insn_mnem(scan_addr)\n if mnem in [\"mul\", \"mult\", \"sll\", \"add\", \"addu\", \"addi\", \"addiu\"]:\n results.append({\n 'risk': 'MEDIUM',\n 'issue': 'Arithmetic (%s) before %s - check for integer overflow' % (mnem, func_name),\n 'detail': {'addr': \"0x%x\" % scan_addr, 'instruction': idc.GetDisasm(scan_addr)}\n })\n break\n scan_count += 1\n scan_addr = idc.prev_head(scan_addr)\n \n return results\n\n\ndef auditFreeCall(func_name):\n \"\"\"\n Audit free() calls for potential double-free vulnerabilities\n \"\"\"\n func_addr = idc.get_name_ea_simple(func_name)\n if func_addr == BADADDR:\n return\n \n print(printFunc(func_name + \" (Double-Free Analysis)\"))\n \n # Group by containing function\n calls_by_func = defaultdict(list)\n \n for call_addr in idautils.CodeRefsTo(func_addr, 0):\n Mnemonics = idc.print_insn_mnem(call_addr)\n if Mnemonics[0:1] == \"j\" or Mnemonics[0:1] == \"b\":\n containing_func = getCallingFunction(call_addr)\n arg_info = getArg(call_addr, 0) # First argument to free\n calls_by_func[containing_func].append({\n 'addr': call_addr,\n 'arg': arg_info\n })\n \n # Analyze each function for potential double-free\n table = PrettyTable([\"function\", \"free_count\", \"addresses\", \"args\", \"risk\"])\n \n for func_start, calls in calls_by_func.items():\n func_name_str = idc.get_func_name(func_start) or (\"0x%x\" % func_start)\n addrs = [(\"0x%x\" % c['addr']) for c in calls]\n args = [c['arg'] for c in calls]\n \n # Risk assessment\n risk = \"LOW\"\n if len(calls) > 1:\n # Check if same argument pattern appears multiple times\n arg_counts = defaultdict(int)\n for arg in args:\n arg_counts[arg] += 1\n \n for arg, count in arg_counts.items():\n if count > 1 and arg != \"get fail\":\n risk = \"HIGH\"\n idc.set_color(calls[0]['addr'], idc.CIC_ITEM, RISK_HIGH)\n break\n else:\n if len(calls) > 2:\n risk = \"MEDIUM\"\n \n table.add_row([func_name_str, len(calls), \", \".join(addrs), \", \".join(args), risk])\n \n print(table)\n\n\ndef auditEnhanced(func_name):\n \"\"\"\n Enhanced audit with additional vulnerability checks\n \"\"\"\n global audit_results\n \n func_addr = idc.get_name_ea_simple(func_name)\n if func_addr == BADADDR:\n return False\n \n print(printFunc(func_name + \" (Enhanced Analysis)\"))\n \n # Determine check type based on function\n check_cmd_injection = func_name in command_execution_function\n check_overflow = func_name in dangerous_functions + attention_function\n check_int_overflow = func_name in memory_alloc_functions\n check_format_string = func_name in format_string_functions\n \n table = PrettyTable([\"addr\", \"risk\", \"issue\", \"detail\"])\n \n for call_addr in idautils.CodeRefsTo(func_addr, 0):\n Mnemonics = idc.print_insn_mnem(call_addr)\n if Mnemonics[0:1] == \"j\" or Mnemonics[0:1] == \"b\":\n findings = []\n \n if check_cmd_injection:\n findings.extend(checkCommandInjection(call_addr, func_name))\n \n if check_overflow:\n findings.extend(checkStackOverflow(call_addr, func_name))\n \n if check_int_overflow:\n findings.extend(checkIntegerOverflow(call_addr, func_name))\n \n if check_format_string:\n findings.extend(checkFormatStringVuln(call_addr, func_name))\n \n # Add findings to table and global results\n for finding in findings:\n risk = finding['risk']\n # Set color based on risk\n if risk == 'HIGH':\n idc.set_color(call_addr, idc.CIC_ITEM, RISK_HIGH)\n elif risk == 'MEDIUM':\n idc.set_color(call_addr, idc.CIC_ITEM, RISK_MEDIUM)\n else:\n idc.set_color(call_addr, idc.CIC_ITEM, RISK_LOW)\n \n detail_str = str(finding.get('detail', ''))[:50]\n table.add_row([\"0x%x\" % call_addr, risk, finding['issue'], detail_str])\n \n # Add to global results for export\n finding['address'] = \"0x%x\" % call_addr\n finding['function'] = func_name\n audit_results.append(finding)\n \n if table.rowcount > 0:\n print(table)\n else:\n print(\" No enhanced findings for %s\" % func_name)\n\n\ndef collectTaintSources():\n \"\"\"\n Collect all external input sources in the binary for taint analysis\n \"\"\"\n global taint_sources\n taint_sources = {}\n \n print(\"\\n[*] Collecting external input sources...\")\n \n for func_name in external_input_functions:\n func_addr = idc.get_name_ea_simple(func_name)\n if func_addr == BADADDR:\n continue\n \n for call_addr in idautils.CodeRefsTo(func_addr, 0):\n Mnemonics = idc.print_insn_mnem(call_addr)\n if Mnemonics[0:1] == \"j\" or Mnemonics[0:1] == \"b\":\n taint_sources[call_addr] = func_name\n \n print(\" Found %d external input call sites\" % len(taint_sources))\n return taint_sources\n\n\n# ============================================================\n# Basic Block Analysis (Fix cross-block issues)\n# ============================================================\n\ndef getBasicBlockBounds(addr):\n \"\"\"Get the basic block boundaries containing the address\"\"\"\n func = idaapi.get_func(addr)\n if not func:\n return (BADADDR, BADADDR)\n \n try:\n flowchart = idaapi.FlowChart(func)\n for block in flowchart:\n if block.start_ea <= addr < block.end_ea:\n return (block.start_ea, block.end_ea)\n except:\n pass\n return (BADADDR, BADADDR)\n\n\ndef getArgAddrInBlock(start_addr, regNum):\n \"\"\"\n Improved argument address detection within basic block boundaries\n Avoids crossing basic block boundaries which could lead to false positives\n \"\"\"\n mipscondition = [\"bn\", \"be\", \"bg\", \"bl\"]\n scan_deep = 50\n count = 0\n reg = \"$a\" + str(regNum)\n \n # Get basic block bounds\n block_start, block_end = getBasicBlockBounds(start_addr)\n \n # try to get in the next (code references from this address)\n for next_addr in idautils.CodeRefsFrom(start_addr, 0):\n if next_addr != BADADDR and reg == idc.print_operand(next_addr, 0):\n return next_addr\n \n # try to get before (within same basic block)\n before_addr = idc.prev_head(start_addr)\n \n while before_addr != BADADDR:\n # Stop if we cross basic block boundary\n if block_start != BADADDR and before_addr < block_start:\n break\n \n if reg == idc.print_operand(before_addr, 0):\n Mnemonics = idc.print_insn_mnem(before_addr)\n if Mnemonics[0:2] in mipscondition:\n pass\n elif Mnemonics[0:1] == \"j\":\n pass\n else:\n return before_addr\n count = count + 1\n if count > scan_deep:\n break \n before_addr = idc.prev_head(before_addr)\n return BADADDR\n\n\n# ============================================================\n# Format String Vulnerability Detection\n# ============================================================\n\ndef checkFormatStringVuln(call_addr, func_name):\n \"\"\"\n Enhanced format string vulnerability detection\n - Detects %n format specifier (write primitive)\n - Detects user-controlled format string\n \"\"\"\n results = []\n \n if func_name not in format_function_offset_dict:\n return results\n \n fmt_arg_idx = format_function_offset_dict[func_name]\n arg_addr = getArgAddrInBlock(call_addr, fmt_arg_idx)\n \n if arg_addr == BADADDR:\n return results\n \n # Try to get the format string\n fmt_info = getFormatString(arg_addr)\n \n if fmt_info != \"get fail\":\n fmt_str = fmt_info[1]\n \n # Check for %n (write primitive - HIGH risk)\n if '%n' in fmt_str or '%hn' in fmt_str or '%hhn' in fmt_str or '%ln' in fmt_str:\n results.append({\n 'risk': 'HIGH',\n 'issue': 'Format string with %%n write primitive',\n 'detail': {'format': fmt_str, 'addr': \"0x%x\" % fmt_info[0]},\n 'type': 'format_string'\n })\n \n # Check for multiple format specifiers (potential overflow)\n fmt_count = len(re.findall(r'%[^%]', fmt_str))\n if fmt_count > 10:\n results.append({\n 'risk': 'MEDIUM',\n 'issue': 'Format string with many specifiers (%d)' % fmt_count,\n 'detail': {'format': fmt_str[:50] + '...', 'count': fmt_count},\n 'type': 'format_string'\n })\n else:\n # Format string is not static - check if user controlled\n source_type, source_info = traceArgSource(call_addr, fmt_arg_idx, depth=15)\n \n if source_type == 'external_input':\n results.append({\n 'risk': 'HIGH',\n 'issue': 'User-controlled format string from %s()' % source_info['function'],\n 'detail': source_info,\n 'type': 'format_string'\n })\n elif source_type == 'stack_var':\n results.append({\n 'risk': 'MEDIUM',\n 'issue': 'Format string from stack variable (needs review)',\n 'detail': source_info,\n 'type': 'format_string'\n })\n \n return results\n\n\n# ============================================================\n# Data Flow Analysis - Forward Tracking\n# ============================================================\n\ndef traceDataFlowForward(start_addr, src_func, max_depth=5):\n \"\"\"\n Trace where data from external input flows to\n Tracks return values from read/recv etc. to dangerous sinks\n \"\"\"\n results = []\n \n if max_depth <= 0:\n return results\n \n # Get the function containing this call\n func = idaapi.get_func(start_addr)\n if not func:\n return results\n \n # For read/recv, the buffer is arg1 ($a1)\n # Track where this buffer is used\n buffer_arg_idx = 1 if src_func in [\"read\", \"recv\", \"recvfrom\", \"fread\"] else 0\n buffer_operand = getArg(start_addr, buffer_arg_idx)\n \n if buffer_operand == \"get fail\":\n return results\n \n # Scan forward in the function to find uses of this buffer\n current_addr = idc.next_head(start_addr)\n func_end = func.end_ea\n scan_count = 0\n max_scan = 100\n \n while current_addr < func_end and current_addr != BADADDR and scan_count < max_scan:\n mnem = idc.print_insn_mnem(current_addr)\n \n # Check if this is a function call\n if mnem in [\"jal\", \"jalr\"]:\n call_target = idc.print_operand(current_addr, 0)\n \n # Check if any dangerous function is called with our tainted buffer\n all_dangerous = dangerous_functions + command_execution_function + format_string_functions\n if call_target in all_dangerous:\n # Check if our buffer is passed as an argument\n for arg_idx in range(4): # MIPS uses $a0-$a3\n arg_operand = getArg(current_addr, arg_idx)\n if buffer_operand in arg_operand or arg_operand in buffer_operand:\n risk = 'HIGH' if call_target in command_execution_function else 'MEDIUM'\n results.append({\n 'risk': risk,\n 'issue': 'Tainted data from %s() flows to %s()' % (src_func, call_target),\n 'detail': {\n 'source': \"0x%x\" % start_addr,\n 'sink': \"0x%x\" % current_addr,\n 'sink_func': call_target,\n 'buffer': buffer_operand\n },\n 'type': 'data_flow'\n })\n break\n \n current_addr = idc.next_head(current_addr)\n scan_count += 1\n \n return results\n\n\ndef analyzeDataFlow():\n \"\"\"\n Perform data flow analysis from external inputs to dangerous sinks\n \"\"\"\n global data_flow_graph\n results = []\n \n print(\"\\n[*] Analyzing data flow from external inputs...\")\n \n for call_addr, func_name in taint_sources.items():\n if func_name in [\"read\", \"recv\", \"recvfrom\", \"fread\", \"fgets\", \"gets\"]:\n flow_results = traceDataFlowForward(call_addr, func_name)\n results.extend(flow_results)\n \n # Store in graph for visualization\n for r in flow_results:\n if 'detail' in r and 'sink' in r['detail']:\n data_flow_graph[call_addr].append((r['detail']['sink'], r['detail']['sink_func']))\n \n print(\" Found %d data flow issues\" % len(results))\n return results\n\n\n# ============================================================\n# Wrapper Function Detection\n# ============================================================\n\ndef detectWrapperFunctions():\n \"\"\"\n Detect wrapper functions that call dangerous functions\n e.g., my_strcpy that internally calls strcpy\n \"\"\"\n global wrapper_functions\n wrapper_functions = {}\n \n print(\"\\n[*] Detecting wrapper functions...\")\n \n all_dangerous = dangerous_functions + command_execution_function\n \n for dangerous_func in all_dangerous:\n func_addr = idc.get_name_ea_simple(dangerous_func)\n if func_addr == BADADDR:\n continue\n \n # Find all callers of this dangerous function\n for call_addr in idautils.CodeRefsTo(func_addr, 0):\n mnem = idc.print_insn_mnem(call_addr)\n if mnem[0:1] != \"j\" and mnem[0:1] != \"b\":\n continue\n \n # Get the function containing this call\n caller_func = idaapi.get_func(call_addr)\n if not caller_func:\n continue\n \n caller_name = idc.get_func_name(caller_func.start_ea)\n if not caller_name:\n continue\n \n # Heuristics for wrapper detection:\n # 1. Function name contains common wrapper patterns\n # 2. Function is small (likely just a wrapper)\n # 3. Function forwards arguments directly\n \n is_wrapper = False\n wrapper_type = None\n \n # Check name patterns\n wrapper_patterns = ['my_', 'safe_', 'wrap_', '_wrapper', '_safe', 'do_', 'internal_']\n for pattern in wrapper_patterns:\n if pattern in caller_name.lower():\n is_wrapper = True\n wrapper_type = 'name_match'\n break\n \n # Check function size (small functions are likely wrappers)\n func_size = caller_func.end_ea - caller_func.start_ea\n if func_size < 100 and not is_wrapper: # Less than ~25 instructions\n # Count how many other calls this function makes\n call_count = 0\n for addr in idautils.FuncItems(caller_func.start_ea):\n m = idc.print_insn_mnem(addr)\n if m in [\"jal\", \"jalr\"]:\n call_count += 1\n \n if call_count <= 2: # Mostly just calls the dangerous function\n is_wrapper = True\n wrapper_type = 'small_func'\n \n if is_wrapper:\n if caller_func.start_ea not in wrapper_functions:\n wrapper_functions[caller_func.start_ea] = {\n 'name': caller_name,\n 'wraps': [],\n 'type': wrapper_type\n }\n wrapper_functions[caller_func.start_ea]['wraps'].append(dangerous_func)\n \n print(\" Found %d potential wrapper functions\" % len(wrapper_functions))\n return wrapper_functions\n\n\ndef auditWrapperFunctions():\n \"\"\"\n Audit detected wrapper functions\n \"\"\"\n if not wrapper_functions:\n detectWrapperFunctions()\n \n if not wrapper_functions:\n print(\" No wrapper functions detected\")\n return\n \n print(printFunc(\"Wrapper Functions\"))\n \n table = PrettyTable([\"wrapper_name\", \"address\", \"wraps\", \"detection_type\"])\n \n for func_addr, info in wrapper_functions.items():\n idc.set_color(func_addr, idc.CIC_FUNC, RISK_MEDIUM)\n table.add_row([\n info['name'],\n \"0x%x\" % func_addr,\n \", \".join(info['wraps']),\n info['type']\n ])\n \n print(table)\n \n # Also audit calls to wrapper functions\n print(\"\\n Calls to wrapper functions:\")\n wrapper_table = PrettyTable([\"wrapper\", \"call_addr\", \"caller_func\"])\n \n for func_addr, info in wrapper_functions.items():\n for call_addr in idautils.CodeRefsTo(func_addr, 0):\n mnem = idc.print_insn_mnem(call_addr)\n if mnem[0:1] == \"j\" or mnem[0:1] == \"b\":\n caller_func = idc.get_func_name(call_addr)\n idc.set_color(call_addr, idc.CIC_ITEM, RISK_MEDIUM)\n wrapper_table.add_row([info['name'], \"0x%x\" % call_addr, caller_func or \"unknown\"])\n \n if wrapper_table.rowcount > 0:\n print(wrapper_table)\n\n\n# ============================================================\n# Result Export Functions\n# ============================================================\n\ndef addFinding(finding):\n \"\"\"Add a finding to the global results list\"\"\"\n global audit_results\n audit_results.append(finding)\n\n\ndef exportResultsJSON(filename):\n \"\"\"Export audit results to JSON file\"\"\"\n try:\n output_path = os.path.join(get_output_dir(), filename)\n with open(output_path, 'w', encoding='utf-8') as f:\n json.dump(audit_results, f, indent=2, ensure_ascii=False)\n print(\"[*] Results exported to: %s\" % output_path)\n return True\n except Exception as e:\n print(\"[!] Error exporting JSON: %s\" % str(e))\n return False\n\n\ndef exportResultsCSV(filename):\n \"\"\"Export audit results to CSV file\"\"\"\n try:\n output_path = os.path.join(get_output_dir(), filename)\n with open(output_path, 'w', newline='', encoding='utf-8') as f:\n if audit_results:\n writer = csv.DictWriter(f, fieldnames=['risk', 'issue', 'type', 'address', 'function', 'detail'])\n writer.writeheader()\n for result in audit_results:\n row = {\n 'risk': result.get('risk', ''),\n 'issue': result.get('issue', ''),\n 'type': result.get('type', ''),\n 'address': result.get('address', ''),\n 'function': result.get('function', ''),\n 'detail': str(result.get('detail', ''))[:100]\n }\n writer.writerow(row)\n print(\"[*] Results exported to: %s\" % output_path)\n return True\n except Exception as e:\n print(\"[!] Error exporting CSV: %s\" % str(e))\n return False\n\n\ndef exportResultsHTML(filename):\n \"\"\"Export audit results to HTML file\"\"\"\n try:\n output_path = os.path.join(get_output_dir(), filename)\n \n html_content = \"\"\"\n\n\n MIPS Audit Report\n \n\n\n

MIPS Security Audit Report

\n
\n Total Findings: %d |\n HIGH: %d |\n MEDIUM: %d |\n LOW: %d\n
\n \n \n \n \n \n \n \n \n \n\"\"\" % (\n len(audit_results),\n len([r for r in audit_results if r.get('risk') == 'HIGH']),\n len([r for r in audit_results if r.get('risk') == 'MEDIUM']),\n len([r for r in audit_results if r.get('risk') == 'LOW'])\n )\n \n for result in audit_results:\n risk = result.get('risk', '')\n html_content += \"\"\" \n \n \n \n \n \n \n \n\"\"\" % (\n risk,\n risk,\n result.get('issue', ''),\n result.get('type', ''),\n result.get('address', ''),\n result.get('function', ''),\n str(result.get('detail', ''))[:100]\n )\n \n html_content += \"\"\"
RiskIssueTypeAddressFunctionDetail
%s%s%s%s%s%s
\n\n\"\"\"\n \n with open(output_path, 'w', encoding='utf-8') as f:\n f.write(html_content)\n print(\"[*] Results exported to: %s\" % output_path)\n return True\n except Exception as e:\n print(\"[!] Error exporting HTML: %s\" % str(e))\n return False\n\n\ndef exportResults(base_filename=\"mipsAudit_results\"):\n \"\"\"Export results as HTML with timestamp\"\"\"\n timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n filename = \"%s_%s.html\" % (base_filename, timestamp)\n exportResultsHTML(filename)\n return filename\n\ndef mipsAudit():\n \"\"\"Main audit function with all enhanced features\"\"\"\n global audit_results, total_functions, processed_functions\n \n # Reset global state\n audit_results = []\n \n # the word create with figlet\n start = '''\n _ _ _ _ _ \n _ __ ___ (_)_ __ ___ / \\ _ _ __| (_) |_ \n| '_ ` _ \\| | '_ \\/ __| / _ \\| | | |/ _` | | __|\n| | | | | | | |_) \\__ \\/ ___ \\ |_| | (_| | | |_ \n|_| |_| |_|_| .__/|___/_/ \\_\\__,_|\\__,_|_|\\__|\n |_| \n code by giantbranch 2018.05\n updated for IDA 7.x+ & Python 3\n enhanced detection v3.0 20260129\n '''\n print(start)\n \n # Load external configuration if available\n load_config()\n \n # Calculate total functions for progress display\n all_functions = (dangerous_functions + attention_function + \n command_execution_function + memory_alloc_functions + \n memory_free_functions + format_string_functions)\n total_functions = len(set(all_functions))\n processed_functions = 0\n \n # Collect taint sources first for enhanced analysis\n collectTaintSources()\n \n print(\"\\n\" + \"=\" * 60)\n print(\" PHASE 1: Basic Function Audit\")\n print(\"=\" * 60)\n \n print(\"\\nAuditing dangerous functions ......\")\n for i, func_name in enumerate(dangerous_functions):\n audit(func_name)\n show_progress(i + 1, len(dangerous_functions), \"Dangerous funcs\")\n \n print(\"\\nAuditing attention function ......\")\n for i, func_name in enumerate(attention_function):\n audit(func_name)\n show_progress(i + 1, len(attention_function), \"Attention funcs\")\n\n print(\"\\nAuditing command execution function ......\")\n for i, func_name in enumerate(command_execution_function):\n audit(func_name)\n show_progress(i + 1, len(command_execution_function), \"Cmd exec funcs\")\n \n print(\"\\n\" + \"=\" * 60)\n print(\" PHASE 2: Enhanced Vulnerability Detection\")\n print(\"=\" * 60)\n \n print(\"\\n[*] Enhanced analysis: Command Injection Detection\")\n for func_name in command_execution_function:\n auditEnhanced(func_name)\n \n print(\"\\n[*] Enhanced analysis: Stack Overflow Detection\")\n for func_name in dangerous_functions:\n auditEnhanced(func_name)\n \n print(\"\\n[*] Enhanced analysis: Integer Overflow Detection\")\n for func_name in memory_alloc_functions:\n auditEnhanced(func_name)\n \n print(\"\\n[*] Enhanced analysis: Format String Vulnerability Detection\")\n for func_name in format_string_functions:\n auditEnhanced(func_name)\n \n print(\"\\n[*] Enhanced analysis: Double-Free Detection\")\n for func_name in memory_free_functions:\n auditFreeCall(func_name)\n \n print(\"\\n\" + \"=\" * 60)\n print(\" PHASE 3: Advanced Analysis\")\n print(\"=\" * 60)\n \n # Data flow analysis\n flow_results = analyzeDataFlow()\n if flow_results:\n print(\"\\nData Flow Analysis Results:\")\n table = PrettyTable([\"risk\", \"issue\", \"source\", \"sink\"])\n for r in flow_results:\n table.add_row([\n r['risk'],\n r['issue'],\n r['detail'].get('source', ''),\n r['detail'].get('sink', '')\n ])\n audit_results.append(r)\n print(table)\n \n # Wrapper function detection\n print(\"\\n[*] Detecting and auditing wrapper functions...\")\n auditWrapperFunctions()\n \n print(\"\\n\" + \"=\" * 60)\n print(\" PHASE 4: Results Summary & Export\")\n print(\"=\" * 60)\n \n # Summary\n high_count = len([r for r in audit_results if r.get('risk') == 'HIGH'])\n medium_count = len([r for r in audit_results if r.get('risk') == 'MEDIUM'])\n low_count = len([r for r in audit_results if r.get('risk') == 'LOW'])\n \n print(\"\\n[*] Audit Summary:\")\n print(\" Total findings: %d\" % len(audit_results))\n print(\" HIGH risk: %d\" % high_count)\n print(\" MEDIUM risk: %d\" % medium_count)\n print(\" LOW risk: %d\" % low_count)\n \n # Export results\n export_file = None\n if audit_results:\n print(\"\\n[*] Exporting results...\")\n export_file = exportResults()\n print(\"\\nExported: %s\" % export_file)\n \n print(\"\\n\" + \"=\" * 60)\n print(\" Finished! Enjoy the result ~\")\n print(\"=\" * 60)\n \n\n# Check processor architecture (may be useful in future)\n# info = idaapi.get_inf_structure()\n#\n# if info.is_64bit():\n# bits = 64\n# elif info.is_32bit():\n# bits = 32\n# else:\n# bits = 16\n#\n# try:\n# is_be = info.is_be()\n# except:\n# is_be = info.mf\n# endian = \"big\" if is_be else \"little\"\n#\n# print('Processor: {}, {}bit, {} endian'.format(info.procName, bits, endian))\n# # Result: Processor: mipsr, 32bit, big endian\n\nmipsAudit()" }, { "path": "oldversion(py2)/mipsAudit.py", "content": "# -*- coding: utf-8 -*-\n\n# reference\n# 《ida pro 权威指南》\n# 《python 灰帽子》\n# 《家用路由器0day漏洞挖掘》\n# https://github.com/wangzery/SearchOverflow/blob/master/SearchOverflow.py\n\nfrom idaapi import *\nfrom prettytable import PrettyTable\n\n\nDEBUG = True\n\n# fgetc,fgets,fread,fprintf,\n# vspritnf\n\n# set function_name\ndangerous_functions = [\n \"strcpy\", \n \"strcat\", \n \"sprintf\",\n \"read\", \n \"getenv\" \n]\n\nattention_function = [\n \"memcpy\",\n \"strncpy\",\n \"sscanf\", \n \"strncat\", \n \"snprintf\",\n \"vprintf\", \n \"printf\"\n]\n\ncommand_execution_function = [\n \"system\", \n \"execve\",\n \"popen\",\n \"unlink\"\n]\n\n# describe arg num of function\n\none_arg_function = [\n \"getenv\",\n \"system\",\n \"unlink\"\n]\n\ntwo_arg_function = [\n \"strcpy\", \n \"strcat\",\n \"popen\"\n]\n\nthree_arg_function = [\n \"strncpy\",\n \"strncat\", \n \"memcpy\",\n \"execve\",\n \"read\"\n]\n\nformat_function_offset_dict = {\n \"sprintf\":1,\n \"sscanf\":1,\n \"snprintf\":2,\n \"vprintf\":0,\n \"printf\":0\n}\n\ndef printFunc(func_name):\n string1 = \"========================================\"\n string2 = \"========== Aduiting \" + func_name + \" \"\n strlen = len(string1) - len(string2)\n return string1 + \"\\n\" + string2 + '=' * strlen + \"\\n\" + string1\n\ndef getFuncAddr(func_name):\n func_addr = LocByName(func_name)\n if func_addr != BADADDR:\n print printFunc(func_name)\n # print func_name + \" Addr : 0x %x\" % func_addr\n return func_addr\n return False\n\ndef getFormatString(addr):\n op_num = 1\n # GetOpType Return value\n #define o_void 0 // No Operand ----------\n #define o_reg 1 // General Register (al, ax, es, ds...) reg\n #define o_mem 2 // Direct Memory Reference (DATA) addr\n #define o_phrase 3 // Memory Ref [Base Reg + Index Reg] phrase\n #define o_displ 4 // Memory Reg [Base Reg + Index Reg + Displacement] phrase+addr\n #define o_imm 5 // Immediate Value value\n #define o_far 6 // Immediate Far Address (CODE) addr\n #define o_near 7 // Immediate Near Address (CODE) addr\n #define o_idpspec0 8 // IDP specific type\n #define o_idpspec1 9 // IDP specific type\n #define o_idpspec2 10 // IDP specific type\n #define o_idpspec3 11 // IDP specific type\n #define o_idpspec4 12 // IDP specific type\n #define o_idpspec5 13 // IDP specific type\n # 如果第二个不是立即数则下一个\n if(GetOpType(addr ,op_num) != 5):\n op_num = op_num + 1\n if GetOpType(addr ,op_num) != 5:\n return \"get fail\"\n op_string = GetOpnd(addr, op_num).split(\" \")[0].split(\"+\")[0].split(\"-\")[0].replace(\"(\", \"\")\n string_addr = LocByName(op_string)\n if string_addr == BADADDR:\n return \"get fail\"\n string = str(GetString(string_addr))\n return [string_addr, string]\n\n\ndef getArgAddr(start_addr, regNum):\n mipscondition = [\"bn\", \"be\" , \"bg\", \"bl\"]\n scan_deep = 50\n count = 0\n reg = \"$a\" + str(regNum)\n # try to get in the next \n next_addr = Rfirst(start_addr)\n if next_addr != BADADDR and reg == GetOpnd(next_addr, 0):\n return next_addr\n # try to get before\n before_addr = RfirstB(start_addr)\n while before_addr != BADADDR:\n if reg == GetOpnd(before_addr, 0):\n Mnemonics = GetMnem(before_addr)\n if Mnemonics[0:2] in mipscondition:\n pass\n elif Mnemonics[0:1] == \"j\":\n pass\n else:\n return before_addr\n count = count + 1\n if count > scan_deep:\n break \n before_addr = RfirstB(before_addr)\n return BADADDR\n\n\ndef getArg(start_addr, regNum):\n mipsmov = [\"move\", \"lw\", \"li\", \"lb\", \"lui\", \"lhu\", \"lbu\", \"la\"]\n arg_addr = getArgAddr(start_addr, regNum)\n if arg_addr != BADADDR:\n Mnemonics = GetMnem(arg_addr) \n if Mnemonics[0:3] == \"add\":\n if GetOpnd(arg_addr, 2) == \"\":\n arg = GetOpnd(arg_addr, 0) + \"+\" + GetOpnd(arg_addr, 1)\n else:\n arg = GetOpnd(arg_addr, 1) + \"+\" + GetOpnd(arg_addr, 2)\n elif Mnemonics[0:3] == \"sub\":\n if GetOpnd(arg_addr, 2) == \"\":\n arg = GetOpnd(arg_addr, 0) + \"-\" + GetOpnd(arg_addr, 1)\n else:\n arg = GetOpnd(arg_addr, 1) + \"-\" + GetOpnd(arg_addr, 2)\n elif Mnemonics in mipsmov:\n arg = GetOpnd(arg_addr, 1) \n else:\n arg = GetDisasm(arg_addr).split(\"#\")[0]\n MakeComm(arg_addr, \"addr: 0x%x \" % start_addr + \"-------> arg\" + str((int(regNum)+1)) + \" : \" + arg)\n return arg\n else:\n return \"get fail\"\n\ndef audit(func_name):\n func_addr = getFuncAddr(func_name) \n if func_addr == False:\n return False\n\n # get arg num and set table\n if func_name in one_arg_function:\n arg_num = 1\n elif func_name in two_arg_function:\n arg_num = 2\n elif func_name in three_arg_function:\n arg_num = 3\n elif func_name in format_function_offset_dict:\n arg_num = format_function_offset_dict[func_name] + 1\n else:\n print \"The %s function didn't write in the describe arg num of function array,please add it to,such as add to `two_arg_function` arary\" % func_name\n return\n # mispcall = [\"jal\", \"jalr\", \"bal\", \"jr\"]\n table_head = [\"func_name\", \"addr\"]\n for num in xrange(0,arg_num):\n table_head.append(\"arg\"+str(num+1))\n if func_name in format_function_offset_dict:\n table_head.append(\"format&value[string_addr, num of '%', fmt_arg...]\")\n table_head.append(\"local_buf_size\")\n table = PrettyTable(table_head)\n\n # get first call\n call_addr = RfirstB(func_addr)\n while call_addr != BADADDR:\n # set color ———— green (red=0x0000ff,blue = 0xff0000)\n SetColor(call_addr, CIC_ITEM, 0x00ff00)\n # set break point\n # AddBpt(call_addr)\n # DelBpt(call_addr)\n\n # if you want to use condition\n # SetBptCnd(ea, 'strstr(GetString(Dword(esp+4),-1, 0), \"SAEXT.DLL\") != -1')\n Mnemonics = GetMnem(call_addr)\n # print \"Mnemonics : %s\" % Mnemonics\n # if Mnemonics in mispcall:\n if Mnemonics[0:1] == \"j\" or Mnemonics[0:1] == \"b\":\n # print func + \" addr : 0x%x\" % call_addr\n if func_name in format_function_offset_dict:\n info = auditFormat(call_addr, func_name, arg_num)\n else:\n info = auditAddr(call_addr, func_name, arg_num)\n table.add_row(info)\n call_addr = RnextB(func_addr, call_addr)\n print table\n # data_addr = DfirstB(func_addr)\n # while data_addr != BADADDR:\n # Mnemonics = GetMnem(data_addr)\n # if DEBUG:\n # print \"Data Mnemonics : %s\" % GetMnem(data_addr)\n # print \"Data addr : 0x %s\" % data_addr\n # data_addr = DnextB(func_addr, data_addr)\n\ndef auditAddr(call_addr, func_name, arg_num):\n addr = \"0x%x\" % call_addr\n ret_list = [func_name, addr]\n # local buf size\n local_buf_size = GetFunctionAttr(call_addr , FUNCATTR_FRSIZE)\n if local_buf_size == BADADDR :\n local_buf_size = \"get fail\"\n else:\n local_buf_size = \"0x%x\" % local_buf_size\n # get arg\n for num in xrange(0,arg_num):\n ret_list.append(getArg(call_addr, num)) \n ret_list.append(local_buf_size)\n return ret_list\n\ndef auditFormat(call_addr, func_name, arg_num):\n addr = \"0x%x\" % call_addr\n ret_list = [func_name, addr]\n # local buf size\n local_buf_size = GetFunctionAttr(call_addr , FUNCATTR_FRSIZE)\n if local_buf_size == BADADDR :\n local_buf_size = \"get fail\"\n else:\n local_buf_size = \"0x%x\" % local_buf_size\n # get arg\n for num in xrange(0,arg_num):\n ret_list.append(getArg(call_addr, num)) \n arg_addr = getArgAddr(call_addr, format_function_offset_dict[func_name])\n string_and_addr = getFormatString(arg_addr)\n format_and_value = []\n if string_and_addr == \"get fail\":\n ret_list.append(\"get fail\")\n else:\n string_addr = \"0x%x\" % string_and_addr[0]\n format_and_value.append(string_addr)\n string = string_and_addr[1]\n fmt_num = string.count(\"%\")\n format_and_value.append(fmt_num)\n # mips arg reg is from a0 to a3\n if fmt_num > 3:\n fmt_num = fmt_num - format_function_offset_dict[func_name] - 1\n for num in xrange(0,fmt_num):\n if arg_num + num > 3:\n break\n format_and_value.append(getArg(call_addr, arg_num + num))\n ret_list.append(format_and_value)\n # format_string = str(getFormatString(arg_addr)[1])\n\n # print \" format String: \" + format_string\n # ret_list.append([string_addr])\n ret_list.append(local_buf_size)\n return ret_list\n\ndef mipsAudit():\n # the word create with figlet\n start = '''\n _ _ _ _ _ \n _ __ ___ (_)_ __ ___ / \\ _ _ __| (_) |_ \n| '_ ` _ \\| | '_ \\/ __| / _ \\| | | |/ _` | | __|\n| | | | | | | |_) \\__ \\/ ___ \\ |_| | (_| | | |_ \n|_| |_| |_|_| .__/|___/_/ \\_\\__,_|\\__,_|_|\\__|\n |_| \n code by giantbranch 2018.05\n '''\n print start\n print \"Auditing dangerous functions ......\"\n for func_name in dangerous_functions:\n audit(func_name)\n \n print \"Auditing attention function ......\"\n for func_name in attention_function:\n audit(func_name)\n\n print \"Auditing command execution function ......\"\n for func_name in command_execution_function:\n audit(func_name)\n \n print \"Finished! Enjoy the result ~\"\n\n# 判断架构的代码，以后或许用得上\n# info = idaapi.get_inf_structure()\n\n# if info.is_64bit():\n# bits = 64\n# elif info.is_32bit():\n# bits = 32\n# else:\n# bits = 16\n\n# try:\n# is_be = info.is_be()\n# except:\n# is_be = info.mf\n# endian = \"big\" if is_be else \"little\"\n\n# print 'Processor: {}, {}bit, {} endian'.format(info.procName, bits, endian)\n# # Result: Processor: mipsr, 32bit, big endian\n\nmipsAudit()" }, { "path": "oldversion(py2)/prettytable.py", "content": "#!/usr/bin/env python\n#\n# Copyright (c) 2009-2013, Luke Maurits \n# All rights reserved.\n# With contributions from:\n# * Chris Clark\n# * Klein Stephane\n#\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n#\n# * Redistributions of source code must retain the above copyright notice,\n# this list of conditions and the following disclaimer.\n# * Redistributions in binary form must reproduce the above copyright notice,\n# this list of conditions and the following disclaimer in the documentation\n# and/or other materials provided with the distribution.\n# * The name of the author may not be used to endorse or promote products\n# derived from this software without specific prior written permission.\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n\n__version__ = \"0.7.2\"\n\nimport copy\nimport csv\nimport random\nimport re\nimport sys\nimport textwrap\nimport itertools\nimport unicodedata\n\npy3k = sys.version_info[0] >= 3\nif py3k:\n unicode = str\n basestring = str\n itermap = map\n iterzip = zip\n uni_chr = chr\n from html.parser import HTMLParser\nelse: \n itermap = itertools.imap\n iterzip = itertools.izip\n uni_chr = unichr\n from HTMLParser import HTMLParser\n\nif py3k and sys.version_info[1] >= 2:\n from html import escape\nelse:\n from cgi import escape\n\n# hrule styles\nFRAME = 0\nALL = 1\nNONE = 2\nHEADER = 3\n\n# Table styles\nDEFAULT = 10\nMSWORD_FRIENDLY = 11\nPLAIN_COLUMNS = 12\nRANDOM = 20\n\n_re = re.compile(\"\\033\\[[0-9;]*m\")\n\ndef _get_size(text):\n lines = text.split(\"\\n\")\n height = len(lines)\n width = max([_str_block_width(line) for line in lines])\n return (width, height)\n \nclass PrettyTable(object):\n\n def __init__(self, field_names=None, **kwargs):\n\n \"\"\"Return a new PrettyTable instance\n\n Arguments:\n\n encoding - Unicode encoding scheme used to decode any encoded input\n field_names - list or tuple of field names\n fields - list or tuple of field names to include in displays\n start - index of first data row to include in output\n end - index of last data row to include in output PLUS ONE (list slice style)\n header - print a header showing field names (True or False)\n header_style - stylisation to apply to field names in header (\"cap\", \"title\", \"upper\", \"lower\" or None)\n border - print a border around the table (True or False)\n hrules - controls printing of horizontal rules after rows. Allowed values: FRAME, HEADER, ALL, NONE\n vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE\n int_format - controls formatting of integer data\n float_format - controls formatting of floating point data\n padding_width - number of spaces on either side of column data (only used if left and right paddings are None)\n left_padding_width - number of spaces on left hand side of column data\n right_padding_width - number of spaces on right hand side of column data\n vertical_char - single character string used to draw vertical lines\n horizontal_char - single character string used to draw horizontal lines\n junction_char - single character string used to draw line junctions\n sortby - name of field to sort rows by\n sort_key - sorting key function, applied to data points before sorting\n valign - default valign for each row (None, \"t\", \"m\" or \"b\")\n reversesort - True or False to sort in descending or ascending order\"\"\"\n\n self.encoding = kwargs.get(\"encoding\", \"UTF-8\")\n\n # Data\n self._field_names = []\n self._align = {}\n self._valign = {}\n self._max_width = {}\n self._rows = []\n if field_names:\n self.field_names = field_names\n else:\n self._widths = []\n\n # Options\n self._options = \"start end fields header border sortby reversesort sort_key attributes format hrules vrules\".split()\n self._options.extend(\"int_format float_format padding_width left_padding_width right_padding_width\".split())\n self._options.extend(\"vertical_char horizontal_char junction_char header_style valign xhtml print_empty\".split())\n for option in self._options:\n if option in kwargs:\n self._validate_option(option, kwargs[option])\n else:\n kwargs[option] = None\n\n self._start = kwargs[\"start\"] or 0\n self._end = kwargs[\"end\"] or None\n self._fields = kwargs[\"fields\"] or None\n\n if kwargs[\"header\"] in (True, False):\n self._header = kwargs[\"header\"]\n else:\n self._header = True\n self._header_style = kwargs[\"header_style\"] or None\n if kwargs[\"border\"] in (True, False):\n self._border = kwargs[\"border\"]\n else:\n self._border = True\n self._hrules = kwargs[\"hrules\"] or FRAME\n self._vrules = kwargs[\"vrules\"] or ALL\n\n self._sortby = kwargs[\"sortby\"] or None\n if kwargs[\"reversesort\"] in (True, False):\n self._reversesort = kwargs[\"reversesort\"]\n else:\n self._reversesort = False\n self._sort_key = kwargs[\"sort_key\"] or (lambda x: x)\n\n self._int_format = kwargs[\"int_format\"] or {}\n self._float_format = kwargs[\"float_format\"] or {}\n self._padding_width = kwargs[\"padding_width\"] or 1\n self._left_padding_width = kwargs[\"left_padding_width\"] or None\n self._right_padding_width = kwargs[\"right_padding_width\"] or None\n\n self._vertical_char = kwargs[\"vertical_char\"] or self._unicode(\"|\")\n self._horizontal_char = kwargs[\"horizontal_char\"] or self._unicode(\"-\")\n self._junction_char = kwargs[\"junction_char\"] or self._unicode(\"+\")\n \n if kwargs[\"print_empty\"] in (True, False):\n self._print_empty = kwargs[\"print_empty\"]\n else:\n self._print_empty = True\n self._format = kwargs[\"format\"] or False\n self._xhtml = kwargs[\"xhtml\"] or False\n self._attributes = kwargs[\"attributes\"] or {}\n \n def _unicode(self, value):\n if not isinstance(value, basestring):\n value = str(value)\n if not isinstance(value, unicode):\n value = unicode(value, self.encoding, \"strict\")\n return value\n\n def _justify(self, text, width, align):\n excess = width - _str_block_width(text)\n if align == \"l\":\n return text + excess * \" \"\n elif align == \"r\":\n return excess * \" \" + text\n else:\n if excess % 2:\n # Uneven padding\n # Put more space on right if text is of odd length...\n if _str_block_width(text) % 2:\n return (excess//2)*\" \" + text + (excess//2 + 1)*\" \"\n # and more space on left if text is of even length\n else:\n return (excess//2 + 1)*\" \" + text + (excess//2)*\" \"\n # Why distribute extra space this way? To match the behaviour of\n # the inbuilt str.center() method.\n else:\n # Equal padding on either side\n return (excess//2)*\" \" + text + (excess//2)*\" \"\n\n def __getattr__(self, name):\n\n if name == \"rowcount\":\n return len(self._rows)\n elif name == \"colcount\":\n if self._field_names:\n return len(self._field_names)\n elif self._rows:\n return len(self._rows[0])\n else:\n return 0\n else:\n raise AttributeError(name)\n \n def __getitem__(self, index):\n\n new = PrettyTable()\n new.field_names = self.field_names\n for attr in self._options:\n setattr(new, \"_\"+attr, getattr(self, \"_\"+attr))\n setattr(new, \"_align\", getattr(self, \"_align\"))\n if isinstance(index, slice):\n for row in self._rows[index]:\n new.add_row(row)\n elif isinstance(index, int):\n new.add_row(self._rows[index])\n else:\n raise Exception(\"Index %s is invalid, must be an integer or slice\" % str(index))\n return new\n\n if py3k:\n def __str__(self):\n return self.__unicode__()\n else:\n def __str__(self):\n return self.__unicode__().encode(self.encoding)\n\n def __unicode__(self):\n return self.get_string()\n\n ##############################\n # ATTRIBUTE VALIDATORS #\n ##############################\n\n # The method _validate_option is all that should be used elsewhere in the code base to validate options.\n # It will call the appropriate validation method for that option. The individual validation methods should\n # never need to be called directly (although nothing bad will happen if they *are*).\n # Validation happens in TWO places.\n # Firstly, in the property setters defined in the ATTRIBUTE MANAGMENT section.\n # Secondly, in the _get_options method, where keyword arguments are mixed with persistent settings\n\n def _validate_option(self, option, val):\n if option in (\"field_names\"):\n self._validate_field_names(val)\n elif option in (\"start\", \"end\", \"max_width\", \"padding_width\", \"left_padding_width\", \"right_padding_width\", \"format\"):\n self._validate_nonnegative_int(option, val)\n elif option in (\"sortby\"):\n self._validate_field_name(option, val)\n elif option in (\"sort_key\"):\n self._validate_function(option, val)\n elif option in (\"hrules\"):\n self._validate_hrules(option, val)\n elif option in (\"vrules\"):\n self._validate_vrules(option, val)\n elif option in (\"fields\"):\n self._validate_all_field_names(option, val)\n elif option in (\"header\", \"border\", \"reversesort\", \"xhtml\", \"print_empty\"):\n self._validate_true_or_false(option, val)\n elif option in (\"header_style\"):\n self._validate_header_style(val)\n elif option in (\"int_format\"):\n self._validate_int_format(option, val)\n elif option in (\"float_format\"):\n self._validate_float_format(option, val)\n elif option in (\"vertical_char\", \"horizontal_char\", \"junction_char\"):\n self._validate_single_char(option, val)\n elif option in (\"attributes\"):\n self._validate_attributes(option, val)\n else:\n raise Exception(\"Unrecognised option: %s!\" % option)\n\n def _validate_field_names(self, val):\n # Check for appropriate length\n if self._field_names:\n try:\n assert len(val) == len(self._field_names)\n except AssertionError:\n raise Exception(\"Field name list has incorrect number of values, (actual) %d!=%d (expected)\" % (len(val), len(self._field_names)))\n if self._rows:\n try:\n assert len(val) == len(self._rows[0])\n except AssertionError:\n raise Exception(\"Field name list has incorrect number of values, (actual) %d!=%d (expected)\" % (len(val), len(self._rows[0])))\n # Check for uniqueness\n try:\n assert len(val) == len(set(val))\n except AssertionError:\n raise Exception(\"Field names must be unique!\")\n\n def _validate_header_style(self, val):\n try:\n assert val in (\"cap\", \"title\", \"upper\", \"lower\", None)\n except AssertionError:\n raise Exception(\"Invalid header style, use cap, title, upper, lower or None!\")\n\n def _validate_align(self, val):\n try:\n assert val in [\"l\",\"c\",\"r\"]\n except AssertionError:\n raise Exception(\"Alignment %s is invalid, use l, c or r!\" % val)\n\n def _validate_valign(self, val):\n try:\n assert val in [\"t\",\"m\",\"b\",None]\n except AssertionError:\n raise Exception(\"Alignment %s is invalid, use t, m, b or None!\" % val)\n\n def _validate_nonnegative_int(self, name, val):\n try:\n assert int(val) >= 0\n except AssertionError:\n raise Exception(\"Invalid value for %s: %s!\" % (name, self._unicode(val)))\n\n def _validate_true_or_false(self, name, val):\n try:\n assert val in (True, False)\n except AssertionError:\n raise Exception(\"Invalid value for %s! Must be True or False.\" % name)\n\n def _validate_int_format(self, name, val):\n if val == \"\":\n return\n try:\n assert type(val) in (str, unicode)\n assert val.isdigit()\n except AssertionError:\n raise Exception(\"Invalid value for %s! Must be an integer format string.\" % name)\n\n def _validate_float_format(self, name, val):\n if val == \"\":\n return\n try:\n assert type(val) in (str, unicode)\n assert \".\" in val\n bits = val.split(\".\")\n assert len(bits) <= 2\n assert bits[0] == \"\" or bits[0].isdigit()\n assert bits[1] == \"\" or bits[1].isdigit()\n except AssertionError:\n raise Exception(\"Invalid value for %s! Must be a float format string.\" % name)\n\n def _validate_function(self, name, val):\n try:\n assert hasattr(val, \"__call__\")\n except AssertionError:\n raise Exception(\"Invalid value for %s! Must be a function.\" % name)\n\n def _validate_hrules(self, name, val):\n try:\n assert val in (ALL, FRAME, HEADER, NONE)\n except AssertionError:\n raise Exception(\"Invalid value for %s! Must be ALL, FRAME, HEADER or NONE.\" % name)\n\n def _validate_vrules(self, name, val):\n try:\n assert val in (ALL, FRAME, NONE)\n except AssertionError:\n raise Exception(\"Invalid value for %s! Must be ALL, FRAME, or NONE.\" % name)\n\n def _validate_field_name(self, name, val):\n try:\n assert (val in self._field_names) or (val is None)\n except AssertionError:\n raise Exception(\"Invalid field name: %s!\" % val)\n\n def _validate_all_field_names(self, name, val):\n try:\n for x in val:\n self._validate_field_name(name, x)\n except AssertionError:\n raise Exception(\"fields must be a sequence of field names!\")\n\n def _validate_single_char(self, name, val):\n try:\n assert _str_block_width(val) == 1\n except AssertionError:\n raise Exception(\"Invalid value for %s! Must be a string of length 1.\" % name)\n\n def _validate_attributes(self, name, val):\n try:\n assert isinstance(val, dict)\n except AssertionError:\n raise Exception(\"attributes must be a dictionary of name/value pairs!\")\n\n ##############################\n # ATTRIBUTE MANAGEMENT #\n ##############################\n\n def _get_field_names(self):\n return self._field_names\n \"\"\"The names of the fields\n\n Arguments:\n\n fields - list or tuple of field names\"\"\"\n def _set_field_names(self, val):\n val = [self._unicode(x) for x in val]\n self._validate_option(\"field_names\", val)\n if self._field_names:\n old_names = self._field_names[:]\n self._field_names = val\n if self._align and old_names:\n for old_name, new_name in zip(old_names, val):\n self._align[new_name] = self._align[old_name]\n for old_name in old_names:\n if old_name not in self._align:\n self._align.pop(old_name)\n else:\n for field in self._field_names:\n self._align[field] = \"c\"\n if self._valign and old_names:\n for old_name, new_name in zip(old_names, val):\n self._valign[new_name] = self._valign[old_name]\n for old_name in old_names:\n if old_name not in self._valign:\n self._valign.pop(old_name)\n else:\n for field in self._field_names:\n self._valign[field] = \"t\"\n field_names = property(_get_field_names, _set_field_names)\n\n def _get_align(self):\n return self._align\n def _set_align(self, val):\n self._validate_align(val)\n for field in self._field_names:\n self._align[field] = val\n align = property(_get_align, _set_align)\n\n def _get_valign(self):\n return self._valign\n def _set_valign(self, val):\n self._validate_valign(val)\n for field in self._field_names:\n self._valign[field] = val\n valign = property(_get_valign, _set_valign)\n\n def _get_max_width(self):\n return self._max_width\n def _set_max_width(self, val):\n self._validate_option(\"max_width\", val)\n for field in self._field_names:\n self._max_width[field] = val\n max_width = property(_get_max_width, _set_max_width)\n \n def _get_fields(self):\n \"\"\"List or tuple of field names to include in displays\n\n Arguments:\n\n fields - list or tuple of field names to include in displays\"\"\"\n return self._fields\n def _set_fields(self, val):\n self._validate_option(\"fields\", val)\n self._fields = val\n fields = property(_get_fields, _set_fields)\n\n def _get_start(self):\n \"\"\"Start index of the range of rows to print\n\n Arguments:\n\n start - index of first data row to include in output\"\"\"\n return self._start\n\n def _set_start(self, val):\n self._validate_option(\"start\", val)\n self._start = val\n start = property(_get_start, _set_start)\n\n def _get_end(self):\n \"\"\"End index of the range of rows to print\n\n Arguments:\n\n end - index of last data row to include in output PLUS ONE (list slice style)\"\"\"\n return self._end\n def _set_end(self, val):\n self._validate_option(\"end\", val)\n self._end = val\n end = property(_get_end, _set_end)\n\n def _get_sortby(self):\n \"\"\"Name of field by which to sort rows\n\n Arguments:\n\n sortby - field name to sort by\"\"\"\n return self._sortby\n def _set_sortby(self, val):\n self._validate_option(\"sortby\", val)\n self._sortby = val\n sortby = property(_get_sortby, _set_sortby)\n\n def _get_reversesort(self):\n \"\"\"Controls direction of sorting (ascending vs descending)\n\n Arguments:\n\n reveresort - set to True to sort by descending order, or False to sort by ascending order\"\"\"\n return self._reversesort\n def _set_reversesort(self, val):\n self._validate_option(\"reversesort\", val)\n self._reversesort = val\n reversesort = property(_get_reversesort, _set_reversesort)\n\n def _get_sort_key(self):\n \"\"\"Sorting key function, applied to data points before sorting\n\n Arguments:\n\n sort_key - a function which takes one argument and returns something to be sorted\"\"\"\n return self._sort_key\n def _set_sort_key(self, val):\n self._validate_option(\"sort_key\", val)\n self._sort_key = val\n sort_key = property(_get_sort_key, _set_sort_key)\n \n def _get_header(self):\n \"\"\"Controls printing of table header with field names\n\n Arguments:\n\n header - print a header showing field names (True or False)\"\"\"\n return self._header\n def _set_header(self, val):\n self._validate_option(\"header\", val)\n self._header = val\n header = property(_get_header, _set_header)\n\n def _get_header_style(self):\n \"\"\"Controls stylisation applied to field names in header\n\n Arguments:\n\n header_style - stylisation to apply to field names in header (\"cap\", \"title\", \"upper\", \"lower\" or None)\"\"\"\n return self._header_style\n def _set_header_style(self, val):\n self._validate_header_style(val)\n self._header_style = val\n header_style = property(_get_header_style, _set_header_style)\n\n def _get_border(self):\n \"\"\"Controls printing of border around table\n\n Arguments:\n\n border - print a border around the table (True or False)\"\"\"\n return self._border\n def _set_border(self, val):\n self._validate_option(\"border\", val)\n self._border = val\n border = property(_get_border, _set_border)\n\n def _get_hrules(self):\n \"\"\"Controls printing of horizontal rules after rows\n\n Arguments:\n\n hrules - horizontal rules style. Allowed values: FRAME, ALL, HEADER, NONE\"\"\"\n return self._hrules\n def _set_hrules(self, val):\n self._validate_option(\"hrules\", val)\n self._hrules = val\n hrules = property(_get_hrules, _set_hrules)\n\n def _get_vrules(self):\n \"\"\"Controls printing of vertical rules between columns\n\n Arguments:\n\n vrules - vertical rules style. Allowed values: FRAME, ALL, NONE\"\"\"\n return self._vrules\n def _set_vrules(self, val):\n self._validate_option(\"vrules\", val)\n self._vrules = val\n vrules = property(_get_vrules, _set_vrules)\n\n def _get_int_format(self):\n \"\"\"Controls formatting of integer data\n Arguments:\n\n int_format - integer format string\"\"\"\n return self._int_format\n def _set_int_format(self, val):\n# self._validate_option(\"int_format\", val)\n for field in self._field_names:\n self._int_format[field] = val\n int_format = property(_get_int_format, _set_int_format)\n\n def _get_float_format(self):\n \"\"\"Controls formatting of floating point data\n Arguments:\n\n float_format - floating point format string\"\"\"\n return self._float_format\n def _set_float_format(self, val):\n# self._validate_option(\"float_format\", val)\n for field in self._field_names:\n self._float_format[field] = val\n float_format = property(_get_float_format, _set_float_format)\n\n def _get_padding_width(self):\n \"\"\"The number of empty spaces between a column's edge and its content\n\n Arguments:\n\n padding_width - number of spaces, must be a positive integer\"\"\"\n return self._padding_width\n def _set_padding_width(self, val):\n self._validate_option(\"padding_width\", val)\n self._padding_width = val\n padding_width = property(_get_padding_width, _set_padding_width)\n\n def _get_left_padding_width(self):\n \"\"\"The number of empty spaces between a column's left edge and its content\n\n Arguments:\n\n left_padding - number of spaces, must be a positive integer\"\"\"\n return self._left_padding_width\n def _set_left_padding_width(self, val):\n self._validate_option(\"left_padding_width\", val)\n self._left_padding_width = val\n left_padding_width = property(_get_left_padding_width, _set_left_padding_width)\n\n def _get_right_padding_width(self):\n \"\"\"The number of empty spaces between a column's right edge and its content\n\n Arguments:\n\n right_padding - number of spaces, must be a positive integer\"\"\"\n return self._right_padding_width\n def _set_right_padding_width(self, val):\n self._validate_option(\"right_padding_width\", val)\n self._right_padding_width = val\n right_padding_width = property(_get_right_padding_width, _set_right_padding_width)\n\n def _get_vertical_char(self):\n \"\"\"The charcter used when printing table borders to draw vertical lines\n\n Arguments:\n\n vertical_char - single character string used to draw vertical lines\"\"\"\n return self._vertical_char\n def _set_vertical_char(self, val):\n val = self._unicode(val)\n self._validate_option(\"vertical_char\", val)\n self._vertical_char = val\n vertical_char = property(_get_vertical_char, _set_vertical_char)\n\n def _get_horizontal_char(self):\n \"\"\"The charcter used when printing table borders to draw horizontal lines\n\n Arguments:\n\n horizontal_char - single character string used to draw horizontal lines\"\"\"\n return self._horizontal_char\n def _set_horizontal_char(self, val):\n val = self._unicode(val)\n self._validate_option(\"horizontal_char\", val)\n self._horizontal_char = val\n horizontal_char = property(_get_horizontal_char, _set_horizontal_char)\n\n def _get_junction_char(self):\n \"\"\"The charcter used when printing table borders to draw line junctions\n\n Arguments:\n\n junction_char - single character string used to draw line junctions\"\"\"\n return self._junction_char\n def _set_junction_char(self, val):\n val = self._unicode(val)\n self._validate_option(\"vertical_char\", val)\n self._junction_char = val\n junction_char = property(_get_junction_char, _set_junction_char)\n\n def _get_format(self):\n \"\"\"Controls whether or not HTML tables are formatted to match styling options\n\n Arguments:\n\n format - True or False\"\"\"\n return self._format\n def _set_format(self, val):\n self._validate_option(\"format\", val)\n self._format = val\n format = property(_get_format, _set_format)\n\n def _get_print_empty(self):\n \"\"\"Controls whether or not empty tables produce a header and frame or just an empty string\n\n Arguments:\n\n print_empty - True or False\"\"\"\n return self._print_empty\n def _set_print_empty(self, val):\n self._validate_option(\"print_empty\", val)\n self._print_empty = val\n print_empty = property(_get_print_empty, _set_print_empty)\n\n def _get_attributes(self):\n \"\"\"A dictionary of HTML attribute name/value pairs to be included in the tag when printing HTML\n\n Arguments:\n\n attributes - dictionary of attributes\"\"\"\n return self._attributes\n def _set_attributes(self, val):\n self._validate_option(\"attributes\", val)\n self._attributes = val\n attributes = property(_get_attributes, _set_attributes)\n\n ##############################\n # OPTION MIXER #\n ##############################\n\n def _get_options(self, kwargs):\n\n options = {}\n for option in self._options:\n if option in kwargs:\n self._validate_option(option, kwargs[option])\n options[option] = kwargs[option]\n else:\n options[option] = getattr(self, \"_\"+option)\n return options\n\n ##############################\n # PRESET STYLE LOGIC #\n ##############################\n\n def set_style(self, style):\n\n if style == DEFAULT:\n self._set_default_style()\n elif style == MSWORD_FRIENDLY:\n self._set_msword_style()\n elif style == PLAIN_COLUMNS:\n self._set_columns_style()\n elif style == RANDOM:\n self._set_random_style()\n else:\n raise Exception(\"Invalid pre-set style!\")\n\n def _set_default_style(self):\n\n self.header = True\n self.border = True\n self._hrules = FRAME\n self._vrules = ALL\n self.padding_width = 1\n self.left_padding_width = 1\n self.right_padding_width = 1\n self.vertical_char = \"|\"\n self.horizontal_char = \"-\"\n self.junction_char = \"+\"\n\n def _set_msword_style(self):\n\n self.header = True\n self.border = True\n self._hrules = NONE\n self.padding_width = 1\n self.left_padding_width = 1\n self.right_padding_width = 1\n self.vertical_char = \"|\"\n\n def _set_columns_style(self):\n\n self.header = True\n self.border = False\n self.padding_width = 1\n self.left_padding_width = 0\n self.right_padding_width = 8\n\n def _set_random_style(self):\n\n # Just for fun!\n self.header = random.choice((True, False))\n self.border = random.choice((True, False))\n self._hrules = random.choice((ALL, FRAME, HEADER, NONE))\n self._vrules = random.choice((ALL, FRAME, NONE))\n self.left_padding_width = random.randint(0,5)\n self.right_padding_width = random.randint(0,5)\n self.vertical_char = random.choice(\"~!@#$%^&*()_+|-=\\{}[];':\\\",./;<>?\")\n self.horizontal_char = random.choice(\"~!@#$%^&*()_+|-=\\{}[];':\\\",./;<>?\")\n self.junction_char = random.choice(\"~!@#$%^&*()_+|-=\\{}[];':\\\",./;<>?\")\n\n ##############################\n # DATA INPUT METHODS #\n ##############################\n\n def add_row(self, row):\n\n \"\"\"Add a row to the table\n\n Arguments:\n\n row - row of data, should be a list with as many elements as the table\n has fields\"\"\"\n\n if self._field_names and len(row) != len(self._field_names):\n raise Exception(\"Row has incorrect number of values, (actual) %d!=%d (expected)\" %(len(row),len(self._field_names)))\n if not self._field_names:\n self.field_names = [(\"Field %d\" % (n+1)) for n in range(0,len(row))]\n self._rows.append(list(row))\n\n def del_row(self, row_index):\n\n \"\"\"Delete a row to the table\n\n Arguments:\n\n row_index - The index of the row you want to delete. Indexing starts at 0.\"\"\"\n\n if row_index > len(self._rows)-1:\n raise Exception(\"Cant delete row at index %d, table only has %d rows!\" % (row_index, len(self._rows)))\n del self._rows[row_index]\n\n def add_column(self, fieldname, column, align=\"c\", valign=\"t\"):\n\n \"\"\"Add a column to the table.\n\n Arguments:\n\n fieldname - name of the field to contain the new column of data\n column - column of data, should be a list with as many elements as the\n table has rows\n align - desired alignment for this column - \"l\" for left, \"c\" for centre and \"r\" for right\n valign - desired vertical alignment for new columns - \"t\" for top, \"m\" for middle and \"b\" for bottom\"\"\"\n\n if len(self._rows) in (0, len(column)):\n self._validate_align(align)\n self._validate_valign(valign)\n self._field_names.append(fieldname)\n self._align[fieldname] = align\n self._valign[fieldname] = valign\n for i in range(0, len(column)):\n if len(self._rows) < i+1:\n self._rows.append([])\n self._rows[i].append(column[i])\n else:\n raise Exception(\"Column length %d does not match number of rows %d!\" % (len(column), len(self._rows)))\n\n def clear_rows(self):\n\n \"\"\"Delete all rows from the table but keep the current field names\"\"\"\n\n self._rows = []\n\n def clear(self):\n\n \"\"\"Delete all rows and field names from the table, maintaining nothing but styling options\"\"\"\n\n self._rows = []\n self._field_names = []\n self._widths = []\n\n ##############################\n # MISC PUBLIC METHODS #\n ##############################\n\n def copy(self):\n return copy.deepcopy(self)\n\n ##############################\n # MISC PRIVATE METHODS #\n ##############################\n\n def _format_value(self, field, value):\n if isinstance(value, int) and field in self._int_format:\n value = self._unicode((\"%%%sd\" % self._int_format[field]) % value)\n elif isinstance(value, float) and field in self._float_format:\n value = self._unicode((\"%%%sf\" % self._float_format[field]) % value)\n return self._unicode(value)\n\n def _compute_widths(self, rows, options):\n if options[\"header\"]:\n widths = [_get_size(field)[0] for field in self._field_names]\n else:\n widths = len(self.field_names) * [0]\n for row in rows:\n for index, value in enumerate(row):\n fieldname = self.field_names[index]\n if fieldname in self.max_width:\n widths[index] = max(widths[index], min(_get_size(value)[0], self.max_width[fieldname]))\n else:\n widths[index] = max(widths[index], _get_size(value)[0])\n self._widths = widths\n\n def _get_padding_widths(self, options):\n\n if options[\"left_padding_width\"] is not None:\n lpad = options[\"left_padding_width\"]\n else:\n lpad = options[\"padding_width\"]\n if options[\"right_padding_width\"] is not None:\n rpad = options[\"right_padding_width\"]\n else:\n rpad = options[\"padding_width\"]\n return lpad, rpad\n\n def _get_rows(self, options):\n \"\"\"Return only those data rows that should be printed, based on slicing and sorting.\n\n Arguments:\n\n options - dictionary of option settings.\"\"\"\n \n # Make a copy of only those rows in the slice range \n rows = copy.deepcopy(self._rows[options[\"start\"]:options[\"end\"]])\n # Sort if necessary\n if options[\"sortby\"]:\n sortindex = self._field_names.index(options[\"sortby\"])\n # Decorate\n rows = [[row[sortindex]]+row for row in rows]\n # Sort\n rows.sort(reverse=options[\"reversesort\"], key=options[\"sort_key\"])\n # Undecorate\n rows = [row[1:] for row in rows]\n return rows\n \n def _format_row(self, row, options):\n return [self._format_value(field, value) for (field, value) in zip(self._field_names, row)]\n\n def _format_rows(self, rows, options):\n return [self._format_row(row, options) for row in rows]\n \n ##############################\n # PLAIN TEXT STRING METHODS #\n ##############################\n\n def get_string(self, **kwargs):\n\n \"\"\"Return string representation of table in current state.\n\n Arguments:\n\n start - index of first data row to include in output\n end - index of last data row to include in output PLUS ONE (list slice style)\n fields - names of fields (columns) to include\n header - print a header showing field names (True or False)\n border - print a border around the table (True or False)\n hrules - controls printing of horizontal rules after rows. Allowed values: ALL, FRAME, HEADER, NONE\n vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE\n int_format - controls formatting of integer data\n float_format - controls formatting of floating point data\n padding_width - number of spaces on either side of column data (only used if left and right paddings are None)\n left_padding_width - number of spaces on left hand side of column data\n right_padding_width - number of spaces on right hand side of column data\n vertical_char - single character string used to draw vertical lines\n horizontal_char - single character string used to draw horizontal lines\n junction_char - single character string used to draw line junctions\n sortby - name of field to sort rows by\n sort_key - sorting key function, applied to data points before sorting\n reversesort - True or False to sort in descending or ascending order\n print empty - if True, stringify just the header for an empty table, if False return an empty string \"\"\"\n\n options = self._get_options(kwargs)\n\n lines = []\n\n # Don't think too hard about an empty table\n # Is this the desired behaviour? Maybe we should still print the header?\n if self.rowcount == 0 and (not options[\"print_empty\"] or not options[\"border\"]):\n return \"\"\n\n # Get the rows we need to print, taking into account slicing, sorting, etc.\n rows = self._get_rows(options)\n\n # Turn all data in all rows into Unicode, formatted as desired\n formatted_rows = self._format_rows(rows, options)\n\n # Compute column widths\n self._compute_widths(formatted_rows, options)\n\n # Add header or top of border\n self._hrule = self._stringify_hrule(options)\n if options[\"header\"]:\n lines.append(self._stringify_header(options))\n elif options[\"border\"] and options[\"hrules\"] in (ALL, FRAME):\n lines.append(self._hrule)\n\n # Add rows\n for row in formatted_rows:\n lines.append(self._stringify_row(row, options))\n\n # Add bottom of border\n if options[\"border\"] and options[\"hrules\"] == FRAME:\n lines.append(self._hrule)\n \n return self._unicode(\"\\n\").join(lines)\n\n def _stringify_hrule(self, options):\n\n if not options[\"border\"]:\n return \"\"\n lpad, rpad = self._get_padding_widths(options)\n if options['vrules'] in (ALL, FRAME):\n bits = [options[\"junction_char\"]]\n else:\n bits = [options[\"horizontal_char\"]]\n # For tables with no data or fieldnames\n if not self._field_names:\n bits.append(options[\"junction_char\"])\n return \"\".join(bits)\n for field, width in zip(self._field_names, self._widths):\n if options[\"fields\"] and field not in options[\"fields\"]:\n continue\n bits.append((width+lpad+rpad)*options[\"horizontal_char\"])\n if options['vrules'] == ALL:\n bits.append(options[\"junction_char\"])\n else:\n bits.append(options[\"horizontal_char\"])\n if options[\"vrules\"] == FRAME:\n bits.pop()\n bits.append(options[\"junction_char\"])\n return \"\".join(bits)\n\n def _stringify_header(self, options):\n\n bits = []\n lpad, rpad = self._get_padding_widths(options)\n if options[\"border\"]:\n if options[\"hrules\"] in (ALL, FRAME):\n bits.append(self._hrule)\n bits.append(\"\\n\")\n if options[\"vrules\"] in (ALL, FRAME):\n bits.append(options[\"vertical_char\"])\n else:\n bits.append(\" \")\n # For tables with no data or field names\n if not self._field_names:\n if options[\"vrules\"] in (ALL, FRAME):\n bits.append(options[\"vertical_char\"])\n else:\n bits.append(\" \")\n for field, width, in zip(self._field_names, self._widths):\n if options[\"fields\"] and field not in options[\"fields\"]:\n continue\n if self._header_style == \"cap\":\n fieldname = field.capitalize()\n elif self._header_style == \"title\":\n fieldname = field.title()\n elif self._header_style == \"upper\":\n fieldname = field.upper()\n elif self._header_style == \"lower\":\n fieldname = field.lower()\n else:\n fieldname = field\n bits.append(\" \" * lpad + self._justify(fieldname, width, self._align[field]) + \" \" * rpad)\n if options[\"border\"]:\n if options[\"vrules\"] == ALL:\n bits.append(options[\"vertical_char\"])\n else:\n bits.append(\" \")\n # If vrules is FRAME, then we just appended a space at the end\n # of the last field, when we really want a vertical character\n if options[\"border\"] and options[\"vrules\"] == FRAME:\n bits.pop()\n bits.append(options[\"vertical_char\"])\n if options[\"border\"] and options[\"hrules\"] != NONE:\n bits.append(\"\\n\")\n bits.append(self._hrule)\n return \"\".join(bits)\n\n def _stringify_row(self, row, options):\n \n for index, field, value, width, in zip(range(0,len(row)), self._field_names, row, self._widths):\n # Enforce max widths\n lines = value.split(\"\\n\")\n new_lines = []\n for line in lines: \n if _str_block_width(line) > width:\n line = textwrap.fill(line, width)\n new_lines.append(line)\n lines = new_lines\n value = \"\\n\".join(lines)\n row[index] = value\n\n row_height = 0\n for c in row:\n h = _get_size(c)[1]\n if h > row_height:\n row_height = h\n\n bits = []\n lpad, rpad = self._get_padding_widths(options)\n for y in range(0, row_height):\n bits.append([])\n if options[\"border\"]:\n if options[\"vrules\"] in (ALL, FRAME):\n bits[y].append(self.vertical_char)\n else:\n bits[y].append(\" \")\n\n for field, value, width, in zip(self._field_names, row, self._widths):\n\n valign = self._valign[field]\n lines = value.split(\"\\n\")\n dHeight = row_height - len(lines)\n if dHeight:\n if valign == \"m\":\n lines = [\"\"] * int(dHeight / 2) + lines + [\"\"] * (dHeight - int(dHeight / 2))\n elif valign == \"b\":\n lines = [\"\"] * dHeight + lines\n else:\n lines = lines + [\"\"] * dHeight\n\n y = 0\n for l in lines:\n if options[\"fields\"] and field not in options[\"fields\"]:\n continue\n\n bits[y].append(\" \" * lpad + self._justify(l, width, self._align[field]) + \" \" * rpad)\n if options[\"border\"]:\n if options[\"vrules\"] == ALL:\n bits[y].append(self.vertical_char)\n else:\n bits[y].append(\" \")\n y += 1\n\n # If vrules is FRAME, then we just appended a space at the end\n # of the last field, when we really want a vertical character\n for y in range(0, row_height):\n if options[\"border\"] and options[\"vrules\"] == FRAME:\n bits[y].pop()\n bits[y].append(options[\"vertical_char\"])\n \n if options[\"border\"] and options[\"hrules\"]== ALL:\n bits[row_height-1].append(\"\\n\")\n bits[row_height-1].append(self._hrule)\n\n for y in range(0, row_height):\n bits[y] = \"\".join(bits[y])\n\n return \"\\n\".join(bits)\n\n ##############################\n # HTML STRING METHODS #\n ##############################\n\n def get_html_string(self, **kwargs):\n\n \"\"\"Return string representation of HTML formatted version of table in current state.\n\n Arguments:\n\n start - index of first data row to include in output\n end - index of last data row to include in output PLUS ONE (list slice style)\n fields - names of fields (columns) to include\n header - print a header showing field names (True or False)\n border - print a border around the table (True or False)\n hrules - controls printing of horizontal rules after rows. Allowed values: ALL, FRAME, HEADER, NONE\n vrules - controls printing of vertical rules between columns. Allowed values: FRAME, ALL, NONE\n int_format - controls formatting of integer data\n float_format - controls formatting of floating point data\n padding_width - number of spaces on either side of column data (only used if left and right paddings are None)\n left_padding_width - number of spaces on left hand side of column data\n right_padding_width - number of spaces on right hand side of column data\n sortby - name of field to sort rows by\n sort_key - sorting key function, applied to data points before sorting\n attributes - dictionary of name/value pairs to include as HTML attributes in the
tag\n xhtml - print
tags if True,
tags if false\"\"\"\n\n options = self._get_options(kwargs)\n\n if options[\"format\"]:\n string = self._get_formatted_html_string(options)\n else:\n string = self._get_simple_html_string(options)\n\n return string\n\n def _get_simple_html_string(self, options):\n\n lines = []\n if options[\"xhtml\"]:\n linebreak = \"
\"\n else:\n linebreak = \"
\"\n\n open_tag = []\n open_tag.append(\"\")\n lines.append(\"\".join(open_tag))\n\n # Headers\n if options[\"header\"]:\n lines.append(\" \")\n for field in self._field_names:\n if options[\"fields\"] and field not in options[\"fields\"]:\n continue\n lines.append(\" \" % escape(field).replace(\"\\n\", linebreak))\n lines.append(\" \")\n\n # Data\n rows = self._get_rows(options)\n formatted_rows = self._format_rows(rows, options)\n for row in formatted_rows:\n lines.append(\" \")\n for field, datum in zip(self._field_names, row):\n if options[\"fields\"] and field not in options[\"fields\"]:\n continue\n lines.append(\" \" % escape(datum).replace(\"\\n\", linebreak))\n lines.append(\" \")\n\n lines.append(\"
%s
%s
\")\n\n return self._unicode(\"\\n\").join(lines)\n\n def _get_formatted_html_string(self, options):\n\n lines = []\n lpad, rpad = self._get_padding_widths(options)\n if options[\"xhtml\"]:\n linebreak = \"
\"\n else:\n linebreak = \"
\"\n\n open_tag = []\n open_tag.append(\"\")\n lines.append(\"\".join(open_tag))\n\n # Headers\n if options[\"header\"]:\n lines.append(\" \")\n for field in self._field_names:\n if options[\"fields\"] and field not in options[\"fields\"]:\n continue\n lines.append(\" %s\" % (lpad, rpad, escape(field).replace(\"\\n\", linebreak)))\n lines.append(\" \")\n\n # Data\n rows = self._get_rows(options)\n formatted_rows = self._format_rows(rows, options)\n aligns = []\n valigns = []\n for field in self._field_names:\n aligns.append({ \"l\" : \"left\", \"r\" : \"right\", \"c\" : \"center\" }[self._align[field]])\n valigns.append({\"t\" : \"top\", \"m\" : \"middle\", \"b\" : \"bottom\"}[self._valign[field]])\n for row in formatted_rows:\n lines.append(\" \")\n for field, datum, align, valign in zip(self._field_names, row, aligns, valigns):\n if options[\"fields\"] and field not in options[\"fields\"]:\n continue\n lines.append(\" %s\" % (lpad, rpad, align, valign, escape(datum).replace(\"\\n\", linebreak)))\n lines.append(\" \")\n lines.append(\"\")\n\n return self._unicode(\"\\n\").join(lines)\n\n##############################\n# UNICODE WIDTH FUNCTIONS #\n##############################\n\ndef _char_block_width(char):\n # Basic Latin, which is probably the most common case\n #if char in xrange(0x0021, 0x007e):\n #if char >= 0x0021 and char <= 0x007e:\n if 0x0021 <= char <= 0x007e:\n return 1\n # Chinese, Japanese, Korean (common)\n if 0x4e00 <= char <= 0x9fff:\n return 2\n # Hangul\n if 0xac00 <= char <= 0xd7af:\n return 2\n # Combining?\n if unicodedata.combining(uni_chr(char)):\n return 0\n # Hiragana and Katakana\n if 0x3040 <= char <= 0x309f or 0x30a0 <= char <= 0x30ff:\n return 2\n # Full-width Latin characters\n if 0xff01 <= char <= 0xff60:\n return 2\n # CJK punctuation\n if 0x3000 <= char <= 0x303e:\n return 2\n # Backspace and delete\n if char in (0x0008, 0x007f):\n return -1\n # Other control characters\n elif char in (0x0000, 0x001f):\n return 0\n # Take a guess\n return 1\n\ndef _str_block_width(val):\n\n return sum(itermap(_char_block_width, itermap(ord, _re.sub(\"\", val))))\n\n##############################\n# TABLE FACTORIES #\n##############################\n\ndef from_csv(fp, field_names = None, **kwargs):\n\n dialect = csv.Sniffer().sniff(fp.read(1024))\n fp.seek(0)\n reader = csv.reader(fp, dialect)\n\n table = PrettyTable(**kwargs)\n if field_names:\n table.field_names = field_names\n else:\n if py3k:\n table.field_names = [x.strip() for x in next(reader)]\n else:\n table.field_names = [x.strip() for x in reader.next()]\n\n for row in reader:\n table.add_row([x.strip() for x in row])\n\n return table\n\ndef from_db_cursor(cursor, **kwargs):\n\n if cursor.description:\n table = PrettyTable(**kwargs)\n table.field_names = [col[0] for col in cursor.description]\n for row in cursor.fetchall():\n table.add_row(row)\n return table\n\nclass TableHandler(HTMLParser):\n\n def __init__(self, **kwargs):\n HTMLParser.__init__(self)\n self.kwargs = kwargs\n self.tables = []\n self.last_row = []\n self.rows = []\n self.max_row_width = 0\n self.active = None\n self.last_content = \"\"\n self.is_last_row_header = False\n\n def handle_starttag(self,tag, attrs):\n self.active = tag\n if tag == \"th\":\n self.is_last_row_header = True\n\n def handle_endtag(self,tag):\n if tag in [\"th\", \"td\"]:\n stripped_content = self.last_content.strip()\n self.last_row.append(stripped_content)\n if tag == \"tr\":\n self.rows.append(\n (self.last_row, self.is_last_row_header))\n self.max_row_width = max(self.max_row_width, len(self.last_row))\n self.last_row = []\n self.is_last_row_header = False\n if tag == \"table\":\n table = self.generate_table(self.rows)\n self.tables.append(table)\n self.rows = []\n self.last_content = \" \"\n self.active = None\n\n\n def handle_data(self, data):\n self.last_content += data\n\n def generate_table(self, rows):\n \"\"\"\n Generates from a list of rows a PrettyTable object.\n \"\"\"\n table = PrettyTable(**self.kwargs)\n for row in self.rows:\n if len(row[0]) < self.max_row_width:\n appends = self.max_row_width - len(row[0])\n for i in range(1,appends):\n row[0].append(\"-\")\n\n if row[1] == True:\n self.make_fields_unique(row[0])\n table.field_names = row[0]\n else:\n table.add_row(row[0])\n return table\n\n def make_fields_unique(self, fields):\n \"\"\"\n iterates over the row and make each field unique\n \"\"\"\n for i in range(0, len(fields)):\n for j in range(i+1, len(fields)):\n if fields[i] == fields[j]:\n fields[j] += \"'\"\n\ndef from_html(html_code, **kwargs):\n \"\"\"\n Generates a list of PrettyTables from a string of HTML code. Each in\n the HTML becomes one PrettyTable object.\n \"\"\"\n\n parser = TableHandler(**kwargs)\n parser.feed(html_code)\n return parser.tables\n\ndef from_html_one(html_code, **kwargs):\n \"\"\"\n Generates a PrettyTables from a string of HTML code which contains only a\n single
\n \"\"\"\n\n tables = from_html(html_code, **kwargs)\n try:\n assert len(tables) == 1\n except AssertionError:\n raise Exception(\"More than one
in provided HTML code! Use from_html instead.\")\n return tables[0]\n\n##############################\n# MAIN (TEST FUNCTION) #\n##############################\n\ndef main():\n\n x = PrettyTable([\"City name\", \"Area\", \"Population\", \"Annual Rainfall\"])\n x.sortby = \"Population\"\n x.reversesort = True\n x.int_format[\"Area\"] = \"04d\"\n x.float_format = \"6.1f\"\n x.align[\"City name\"] = \"l\" # Left align city names\n x.add_row([\"Adelaide\", 1295, 1158259, 600.5])\n x.add_row([\"Brisbane\", 5905, 1857594, 1146.4])\n x.add_row([\"Darwin\", 112, 120900, 1714.7])\n x.add_row([\"Hobart\", 1357, 205556, 619.5])\n x.add_row([\"Sydney\", 2058, 4336374, 1214.8])\n x.add_row([\"Melbourne\", 1566, 3806092, 646.9])\n x.add_row([\"Perth\", 5386, 1554769, 869.4])\n print(x)\n \nif __name__ == \"__main__\":\n main()\n" } ]