Pyotp实现命令行查看2FA

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

#!/usr/bin/env python3
import time
import sys
from datetime import datetime

import pyotp
from wasabi import msg


"""
accounts 为账户列表，支持两种格式：
1. name、issuer、totpSecret 必填
2. 使用uri: url 必填(需要包含name和issuer)
"""

accounts = [
    {
        "totpSecret": "xxxxxxxxxxxxxxxxx", # 密钥
        "isShow": False, # 是否显示
        "name": "myaccount", # 名称
        "issuer": "myissuer", # 发行者
        "algorithm": "ALGO_SHA1", # 保留字段, 可空
        "secret": "xxxxx", # 保留字段, 可空
        "digits": 1, # 保留字段, 可空
        "type": "OTP_TOTP", # 保留字段, 可空

    },
    {
        "uri": "otpauth://totp/名称?secret=xxxxxxxxxxxxxxxxx&issuer=发布者",
    }
]


def get_code(totp_secret: str):
    totp = pyotp.TOTP(totp_secret)
    return totp.now()


def echo_code(c: str):
    sec = datetime.now().second
    if 0 <= sec < 20 or 30 <= sec < 50:
        msg.good(c)
    elif 20 <= sec < 25 or 50 <= sec < 55:
        msg.warn(c)
    else:
        msg.fail(c)

    print()


def main():
    for account in accounts:
        is_show = account.get("isShow")
        if not (is_show or is_show is None):
            continue

        uri = account.get("uri")
        if uri:
            totp_obj = pyotp.parse_uri(uri)
            totp_secret = totp_obj.secret
            name = totp_obj.name
            issuer = totp_obj.issuer
        else:
            totp_secret = account["totpSecret"]
            name = account.get("name")
            issuer = account.get("issuer")

        title = "%s:%s" % (issuer, name)
        if len(sys.argv) == 1:
            print(title)
            echo_code(get_code(totp_secret))

        elif title.lower().find(sys.argv[1].lower()) != -1:
            print(title)
            echo_code(get_code(totp_secret))


if __name__ == '__main__':
    main()