python趣味题-数字加密

# coding:utf-8

# 实现对一个四位整数的加密过程，返回加密数：
# 对该四位数的每一位都加5，并用和对10取余代替该数字，
# 在将第一位与第四位交换，第二位与第三位交换，得到一个加密数


def encrypt(num):
    """
    加密一个四位数字的算法
    """
    
    s = str(num)
    a = (int(s[3]) + 5) % 10
    b = (int(s[1]) + 5) % 10
    c = (int(s[2]) + 5) % 10
    d = (int(s[0]) + 5) % 10

    return a * 1000 + b * 100 + c * 10 + d

def test_encrypt():
    numbers = [
            (1234, 9786),
            (8399, 4843),
            (1535, 86),
            ]
    for n in numbers:
        assert encrypt(n[0]) == n[1]