為ssh伺服器新增2fa認證，一個python指令碼全搞定

janbar發表於2024-06-28
伺服器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))
  1. 修改/etc/passwd檔案，將你希望登陸的使用者的預設shell改為/bin/login
# 如下所示修改了root的預設shell
vim /etc/passwd
root:x:0:0:root:/root:/bin/login
  1. 然後重新登陸ssh，此時需要輸入2fa驗證碼才能成功。上面程式碼做了3次錯誤輸入錯誤自動退出ssh登陸狀態，超過60秒未輸入任何字元也自動退出ssh登陸狀態。注意到verify()方法，是為了ssh登陸後相同公網ip客戶端登陸ssh時不需要重複輸入2fa驗證碼，我這樣做是為了方便vscode遠端或scp等其他不方便輸入驗證碼的客戶端免密登陸。當然伺服器判斷沒有任何該公網ip客戶端時需要輸入驗證碼。需要注意這行程式碼：sys.argv[0] = 'bash'，表示成功輸入驗證碼後開啟的shell，可自行修改。

相關文章