有没有办法防止这些哈希函数中类似的代码块重复？

问题描述 投票：0回答：4

我创建了这个程序来计算给定文件的 sha256 或 sha512 哈希值并将计算摘要为十六进制。

它由 5 个文件组成，其中 4 个是自定义模块，1 个是主模块。

我在不同的模块中有两个函数，但这些函数的唯一区别是一个变量。见下图：

来自sha256.py

def get_hash_sha256():
    global sha256_hash
    filename = input("Enter the file name: ")
    sha256_hash = hashlib.sha256()
    with open(filename, "rb") as f:
        for byte_block in iter(lambda: f.read(4096),b""):
            sha256_hash.update(byte_block)
#       print("sha256 valule: \n" + Color.GREEN + sha256_hash.hexdigest())
        print(Color.DARKCYAN + "sha256 value has been calculated")
        color_reset()

来自sha512.py

def get_hash_sha512():
    global sha512_hash
    filename = input("Enter the file name: ")
    sha512_hash = hashlib.sha512()
    with open(filename, "rb") as f:
        for byte_block in iter(lambda: f.read(4096),b""):
            sha512_hash.update(byte_block)
#       print("sha512 valule: \n" + Color.GREEN + sha512_hash.hexdigest())
        print(Color.DARKCYAN + "sha512 value has been calculated")
        color_reset()

这些函数在我的 simple_sha_find.py 文件中调用：

def which_hash():
    sha256_or_sha512 = input("Which hash do you want to calculate: sha256 or sha512? \n")
    if sha256_or_sha512 == "sha256":
        get_hash_sha256()
        verify_checksum_sha256()
    elif sha256_or_sha512 == "sha512":
        get_hash_sha512()
        verify_checksum_sha512()
    else:
        print("Type either sha256 or sha512. If you type anything else the program will close...like this.")
        sys.exit()

if __name__ == "__main__":
    which_hash()

如您所见，将调用的函数基于用户输入。如果用户输入 sha256，则会触发 sha256.py 中的函数，但如果用户输入 sha512，则会触发 sha512.py 中的函数

该应用程序可以工作，但我知道我可以减少它的冗余，但我不知道如何做。

如何定义一次 get_hash_sha---() 和 verify_checksum_sha---() 函数，并根据用户选择 sha256 还是 sha512 执行适当的计算？

我对该程序进行了一些编码变体。

我已将其创建为一个文件，并创建不同的模块并从这些模块调用函数。

无论哪种情况，我都会重复，但我知道这往往会违背自动化的目的。

python sha
4个回答
2
投票

您可以将这两个函数合并为一个函数：

import hashlib

def get_hash(hash_type):
    if hash_type == 'sha256':
        hash_obj= hashlib.sha256()
    elif hash_type == 'sha512':
        hash_obj = hashlib.sha512()
    else:
        print("Invalid hash type.Please choose 'sha256'or'sha512'")
        return

    filename = input("Enter the fileename:  ")
    try:
        with open(filename,"rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                hash_obj.update(byte_block)
        print(Color.DARKCYAN + f"{hash_type} value has been calculated")
        color_reset()
    except FileNotFoundError:
        print(f"File '{filename}' not found.")

def which_hash():
    sha_type =input("Which hash do you want to calculate: sha256 or sha512? \n").lower()
    if sha_type in ['sha256', 'sha512']:
        get_hash(sha_type)
        verify_checksum(sha_type)
    else:
        print("Type sha256 or sha512. If you type anything else program will close. .")
        sys.exit()

if __name__ == "__main__":
    which_hash()

使用 Enum 而不是纯文本也是最佳实践：

from enum import Enum

class HashType(Enum):
    SHA256 = 'sha256'
    SHA512 = 'sha512'

所以你可以改变

if hash_type == HashType.SHA256:
    hash_obj = hashlib.sha256()
elif hash_type == HashType.SHA512:
    hash_obj = hashlib.sha512()

def which_hash():
    sha_type_input = input("Which hash do you want to calculate: sha256 or sha512? \n").lower()
    
    try:
        sha_type = HashType(sha_type_input)
        get_hash(sha_type)
        verify_checksum(sha_type)
    except ValueError:
        print("Type either sha256 or sha512. If you type anything else the program will close.")
        sys.exit()
2
投票

您可以通过传递相关哈希函数作为参数来概括生成哈希的函数。

类似这样的：

from hashlib import sha256, sha512
from typing import Callable

HASH_MAP: dict[str, Callable] = {"sha256": sha256, "sha512": sha512}
CHUNK = 4096

def make_hash(filename: str, hash_function: Callable) -> str:
    hf = hash_function()
    with open(filename, "rb") as data:
        while buffer := data.read(CHUNK):
            hf.update(buffer)
    return hf.hexdigest()

def main():
    filename = input("Enter filename: ")
    func = input(f"Enter hash type {tuple(HASH_MAP)}: ")
    if hfunc := HASH_MAP.get(func):
        print(make_hash(filename, hfunc))
    else:
        print("Invalid hash type selection")

if __name__ == "__main__":
    main()

如果您随后想要添加更多哈希算法，您只需适当编辑 HASH_MAP 字典即可。不需要更改其他代码

0
投票

您可以重构函数，使哈希类型成为参数。可能还避免使用全局变量，并将任何交互式 I/O 留给调用代码。

我还更改了代码以在出现问题时引发错误。对于非常简单的程序来说，仅仅打印一条错误消息就可以了，但是可重用的代码需要正确区分成功和失败。

def get_hash(hash_type, filename):
    if hash_type == 'sha256':
        hash_obj = hashlib.sha256()
    elif hash_type == 'sha512':
        hash_obj = hashlib.sha512()
    else:
        raise ValueError("Invalid hash type.Please choose 'sha256'or'sha512'")

    # Don't trap the error
    with open(filename,"rb") as f:
        for byte_block in iter(lambda: f.read(4096), b""):
                hash_obj.update(byte_block)

def which_hash():
    sha_type = input("Which hash do you want to calculate: sha256 or sha512? \n").lower()
    if sha_type in ['sha256', 'sha512']:
        filename = input("File name: ")
        hash = get_hash(sha_type, filename)

        print(f"{Color.DARKCYAN}{hash_type} value has been calculated")
        color_reset()

        verify_checksum(hash, sha_type, filename)
    else:
        raise ValueError("Type sha256 or sha512")

这种“案例和调用”正是面向对象编程旨在避免的，但也许在您的编程之旅中解决该主题还为时过早。

0
投票

您可以将算法名称作为字符串提供给 

hashlib.file_digest
。

import hashlib

options = 'sha256', 'sha512'

# Choose algorithm
opts = ' or '.join(options)
alg = input(f"Which hash do you want to calculate: {opts}? \n")
if alg not in options:
    print(f"Type either {opts}. If you type anything else the program will close...like this.")
    sys.exit()

# Choose file and hash it
filename = input("Enter the file name: ")
with open(filename, "rb") as f:
    digest = hashlib.file_digest(f, alg)

print(f"{alg} value has been calculated")

在线尝试！

最新问题
© www.soinside.com 2019 - 2024. All rights reserved.