Oct 26, 2018

python 实现文件校验

今天又学习了点 argparse 库的东西——子命令实现的正确姿势(常规操作)，借着实现一个文件校验、比较的小工具来巩固一下刚学到的子命令实现。

本文实现了一个名为 hashpare^[之所以叫 hashpare, 是因为工具使用 hash 值对文件进行比较，就是 hash compare，合称 hashpare] 的工具，我已经托管到 github 上tools-with-script/hashpare。总感觉直接上代码有点粗暴，下面分步骤讲解实现过程。

脚本准备

新建文件： touch hashpare
添加文件运行权限： chmod +x hashpare
注明脚本解释器，文中首行添加：
```
#!/usr/bin/env python3
```
本文会用到三个内建库，分别是 os、hashlib、argparse，导入它们：
```
import os
import hashlib
import argparse
```
- os: 用来操作文件或目录的路径
- hashlib: 用来计算文件的 hash 值
- argparse: 用来解析命令行参数，本文会用到 add_subparsers 添加子解析器集，从而添加子命令^[Many programs split up their functionality into a number of sub-commands, for example, the svn program can invoke sub-commands like svn checkout, svn update, and svn commit…ArgumentParser supports the creation of such sub-commands with the add_subparsers() method…, Sub-commands]

计算校验

这里使用 hashlib 来计算文件的校验码，封装函数如下：

def calculate_hashvalue(file, algorithm='sha1'):
    """计算文件校验值，可选校验算法 sha1 或 md5
    """
    if algorithm not in ['sha1', 'md5']:
        print(f'[-] 不支持校验算法: {algorithm} [-]')
        return (algorithm, None)
    hashobj = hashlib.sha1() if algorithm == 'sha1' else hashlib.md5()
    with open(file, 'rb') as f:
        hashobj.update(f.read())
    return algorithm, hashobj.hexdigest()

方法内部实现很简单，就是根据指定的校验算法计算指定文件的校验码，并返回校验算法名称和校验码。

解析参数

将解析参数相关配置和操作封装到 process_with_args 函数中：

def process_with_args():
    """解析命令行参数
    """
    parser = argparse.ArgumentParser(description='计算文件校验，检查校验值，比较两个文件')
    commands = parser.add_subparsers(help='子命令')
    check = commands.add_parser('check', help='检查校验值命令')
    check.add_argument('file', help='指明文件路径')
    check.add_argument('hash', help='指定校验码')
    check.add_argument('-a', '--algorithm', default='sha1', choices=['sha1', 'md5'], help='指定校验算法')
    check.set_defaults(func=check_hash)
    compare = commands.add_parser('compare', help='比较两个文件')
    compare.add_argument('file', nargs=2, help='指定要比较的两个文件')
    compare.add_argument('-a', '--algorithm', default='sha1', choices=['sha1', 'md5'], help='指定校验算法')
    compare.set_defaults(func=compare_hash)
    calculate = commands.add_parser('calculate', help='计算校验码')
    calculate.add_argument('file', nargs='+', help='指定要计算校验的文件')
    calculate.add_argument('-a', '--algorithm', default='sha1', choices=['sha1', 'md5'], help='指定校验算法')
    calculate.set_defaults(func=calculate_hash)
    return parser.parse_args()

上述代码很想可以看到，有三个子命令：

check: 主要是用来检验文件的校验值是不是与给定的校验码一致
compare: 比较两个文件的校验值从而验证内容是否一致
calculate: 计算给定文件的校验码

另外，每个子命令的解析器都绑定了一个处理函数，这种方式是处理子命令更加方便。

check子命令处理

check 子命令绑定了方法 check_hash:

def check_hash(args):
    args.file = os.path.abspath(args.file)
    algorithm, hashvalue = calculate_hashvalue(args.file, args.algorithm)
    print(f'\n    {args.file}\n    {algorithm}: {hashvalue}')
    result = '一致' if hashvalue == args.hash else '不一致'
    print(f'\n    文件校验码与 {args.hash} {result}\n')

这个方法会打印文件的名称和校验码，最终打印与给定校验码的比较结果。

compare子命令处理

compare 子命令绑定了方法 compare_hash:

def compare_hash(args):
    print()
    hashvalue = [None, None]
    for i, filepath in enumerate(args.file):
        filepath = os.path.abspath(filepath)
        algorithm, hashvalue[i] = calculate_hashvalue(filepath, args.algorithm)
        print(f'    name: {filepath}\n    {algorithm}: {hashvalue[i]}\n')
    result = '一致' if hashvalue[0] == hashvalue[1] else '不一致'
    print(f'    两文件内容{result}\n')

这里最终会打印两个文件的名称、校验码，以及最终的比较结果。

calculate子命令处理

calculate 子命令绑定了方法calculate_hash:

def calculate_hash(args):
    print()
    for filepath in args.file:
        filepath = os.path.abspath(filepath)
        algorithm, hashvalue = calculate_hashvalue(filepath, args.algorithm)
        print(f'    name: {filepath}\n    {algorithm}: {hashvalue}\n')

最终会打印所有指定文件的名称及计算的校验码。

最终实现

最终实现非常简单，只需要如下即可：

if __name__ == '__main__':
    args = process_with_args()
    args.func(args)

测试

工具的帮助信息如下

$ ./hashpare -h
usage: hashpare [-h] {check,compare,calculate} ...

计算文件校验，检查校验值，比较两个文件

positional arguments:
  {check,compare,calculate}
                        子命令
    check               检查校验值命令
    compare             比较两个文件
    calculate           计算校验码

optional arguments:
  -h, --help            show this help message and exit

测试check子命令

这里针对 README.md 文件进行测试

$ ./hashpare check README.md 1234567890

    /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    sha1: daf0fe7a2fab819444d21b1c082e161f6d613875

    文件校验码与 1234567890 不一致

$ ./hashpare check README.md daf0fe7a2fab819444d21b1c082e161f6d613875 -a sha1

    /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    sha1: daf0fe7a2fab819444d21b1c082e161f6d613875

    文件校验码与 daf0fe7a2fab819444d21b1c082e161f6d613875 一致

$ ./hashpare check README.md daf0fe7a2fab819444d21b1c082e161f6d613875 -a md5

    /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    md5: 0ea8599b24a76f586a6610bf9156b1e7

    文件校验码与 daf0fe7a2fab819444d21b1c082e161f6d613875 不一致

测试compare子命令

针对文件 README.md 测试

$ ./hashpare compare README.md README.md

    name: /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    sha1: daf0fe7a2fab819444d21b1c082e161f6d613875

    name: /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    sha1: daf0fe7a2fab819444d21b1c082e161f6d613875

    两文件内容一致

$ ./hashpare compare README.md README.md -a md5

    name: /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    md5: 0ea8599b24a76f586a6610bf9156b1e7

    name: /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    md5: 0ea8599b24a76f586a6610bf9156b1e7

    两文件内容一致

测试calculate子命令

这里计算 README.md 的校验码：

$ ./hashpare calculate README.md

    name: /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    sha1: daf0fe7a2fab819444d21b1c082e161f6d613875

$ ./hashpare calculate README.md -a md5

    name: /Users/5km/Documents/workspace/python/tools-with-script/hashpare/README.md
    md5: 0ea8599b24a76f586a6610bf9156b1e7