为ssh服务器添加2fa认证，一个python脚本全搞定

服务器ssh如果被别人登陆就是一场灾难，所以我研究了ssh认证，我发现Google Authenticator PAM可以实现ssh的2fa认证，但是安装和配置比较麻烦。因此我用python实现了ssh的2fa认证。考虑到很多Linux服务器默认安装python，所以我用py脚本，并只使用标准库，不需要安装第三方py库，方便部署。

1. 首先保存如下脚本到文件：/bin/login，设置执行权限：chmod +x /bin/login，记得修改TOTP_SECRET密钥

#!/bin/env python
import os
import sys
import signal
import getpass
import subprocess
import hmac
import time
import base64
import hashlib

# 随机生成长度为16的全大写字符串作为2fa的密钥
TOTP_SECRET = 'KHGSRAEPVAFPPAGX'

try:
    def gen_totp(secret: str, input=int(time.time()/30), digits=6):
        if (missing_padding := len(secret) % 8) != 0:
            secret += "=" * (8 - missing_padding)
        byte_secret = base64.b32decode(secret, casefold=True)

        result = bytearray()
        while input != 0:
            result.append(input & 0xFF)
            input >>= 8
        byte_input = bytes(bytearray(reversed(result)).rjust(8, b"\0"))

        hasher = hmac.new(byte_secret, byte_input, hashlib.sha1)
        hmac_hash = bytearray(hasher.digest())
        offset = hmac_hash[-1] & 0xF
        code = ((hmac_hash[offset] & 0x7F) << 24
                | (hmac_hash[offset + 1] & 0xFF) << 16
                | (hmac_hash[offset + 2] & 0xFF) << 8
                | (hmac_hash[offset + 3] & 0xFF))
        str_code = str(10_000_000_000 + (code % 10**digits))
        return str_code[-digits:]

    def read_totp_code():
        def timeout(signum, frame): raise TimeoutError
        signal.signal(signal.SIGALRM, timeout)
        signal.alarm(60)

        flag = 0  # no
        try:
            if getpass.getpass('code: ') == gen_totp(TOTP_SECRET):
                flag = 1  # yes
        except BaseException:
            flag = 2  # timeout,ctrl+c

        signal.alarm(0)
        return flag

    def verify():
        if len(sshClient := os.getenv('SSH_CLIENT', '').split()) != 3:
            return True
        user = os.getenv('USER', '')
        tty = os.getenv('SSH_TTY', '').lstrip('/dev/')

        with subprocess.Popen('who', stdout=subprocess.PIPE) as who:
            while line := who.stdout.readline():
                line = line.decode()
                if user in line and sshClient[0] in line and (tty == '' or tty not in line):
                    return False
        return True

    def main():
        if verify():
            flag = 0
            for _ in range(3):
                if (flag := read_totp_code()) > 0:
                    break
                print('Login incorrect')
            if flag != 1:
                return
        sys.argv[0] = 'bash'
        subprocess.call(sys.argv, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr)

    main()
except BaseException as e:
    base = os.path.dirname(os.path.abspath(__file__))
    with open(os.path.join(base, 'login.log'), 'w') as fw:
        fw.write(str(e))

2. 修改/etc/passwd文件，将你希望登陆的用户的默认shell改为/bin/login

# 如下所示修改了root的默认shell
vim /etc/passwd
root:x:0:0:root:/root:/bin/login

3. 然后重新登陆ssh，此时需要输入2fa验证码才能成功。上面代码做了3次错误输入错误自动退出ssh登陆状态，超过60秒未输入任何字符也自动退出ssh登陆状态。注意到verify()方法，是为了ssh登陆后相同公网ip客户端登陆ssh时不需要重复输入2fa验证码，我这样做是为了方便vscode远程或scp等其他不方便输入验证码的客户端免密登陆。当然服务器判断没有任何该公网ip客户端时需要输入验证码。需要注意这行代码：sys.argv[0] = 'bash'，表示成功输入验证码后打开的shell，可自行修改。