164 lines
4.3 KiB
Python
Executable File
164 lines
4.3 KiB
Python
Executable File
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
from __future__ import print_function
|
|
import argparse
|
|
import base64
|
|
import ConfigParser
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import os
|
|
import struct
|
|
import time
|
|
|
|
|
|
|
|
def die(reason):
|
|
# Terminate with an error message
|
|
print("Ecountered an error, terminating")
|
|
print("Error message -", reason)
|
|
exit(1)
|
|
|
|
|
|
def get_arguments():
|
|
# Get input from the command line
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("site", nargs="*")
|
|
default_cfg_path = os.environ['HOME'] + '/.totprc'
|
|
parser.add_argument("-c", "--config", default=default_cfg_path, help='Path to config-file. Defaults to ~/.totprc')
|
|
parser.add_argument("-a", "--andOTP", help='Path to andOTP backup file')
|
|
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def menu(data):
|
|
# Print a pretty menu to choose from
|
|
keynum = 1
|
|
print(" ---------------- Available keys ----------------")
|
|
print(" |")
|
|
for p in data:
|
|
print(" | %2d - %s" % (keynum, p["label"]))
|
|
keynum = keynum + 1
|
|
|
|
print(" |")
|
|
print(" ------------------------------------------------")
|
|
print(" | Please select which key to generate [1 - %d] |" % keynum)
|
|
print(" ------------------------------------------------")
|
|
print()
|
|
return keynum
|
|
|
|
|
|
def print_OTP(secret):
|
|
# Generate the key and pretty print it
|
|
value = TOTP(secret).generate()
|
|
# Format response like XXX XXX
|
|
print(value[:3], value[3:])
|
|
|
|
|
|
def read_config(args):
|
|
# Read from the config file
|
|
config = ConfigParser.ConfigParser()
|
|
|
|
try:
|
|
config.read(args.config)
|
|
except:
|
|
# Error in config file
|
|
die("Could not read %s" % args.config)
|
|
|
|
# Verify that the config is correct
|
|
try:
|
|
config.get('totp', 'andOTPfile')
|
|
except:
|
|
die("Could not find path to 'andOTPfile' in %s" % args.config)
|
|
|
|
if not os.path.isfile(config.get('totp', 'andOTPfile')):
|
|
die("The file %s does not exist" % config.get('totp', 'andOTPfile'))
|
|
|
|
return config.get('totp', 'andOTPfile')
|
|
|
|
|
|
def read_file(andOTPfile):
|
|
# Open and parse the data file
|
|
try:
|
|
with open(andOTPfile) as filen:
|
|
data = json.load(filen)
|
|
except:
|
|
die("Error parsing JSON, corrupt andOTP-file?")
|
|
|
|
return data
|
|
|
|
|
|
class TOTP():
|
|
class TOTPException(Exception):
|
|
def __init__(self, msg):
|
|
self.msg = msg
|
|
|
|
def __str__(self):
|
|
return repr(self.msg)
|
|
|
|
|
|
def __init__(self, secret):
|
|
if not secret:
|
|
raise TOTP.TOTPException('Invalid secret')
|
|
self.secret = secret
|
|
|
|
def generate(self):
|
|
try:
|
|
key = base64.b32decode(self.secret)
|
|
num = int(time.time()) // 30
|
|
msg = struct.pack('>Q', num)
|
|
|
|
# Take a SHA1 HMAC of key and binary-packed time value
|
|
digest = hmac.new(key, msg, hashlib.sha1).digest()
|
|
|
|
# Last 4 bits of the digest tells which 4 bytes to use
|
|
offset = ord(digest[19]) & 15
|
|
token_base = digest[offset : offset+4]
|
|
|
|
# Unpack that into an integer and strip it down
|
|
token_val = struct.unpack('>I', token_base)[0] & 0x7fffffff
|
|
token_num = token_val % 1000000
|
|
|
|
# Pad with leading zeroes
|
|
token = '{0:06d}'.format(token_num)
|
|
|
|
return token
|
|
except:
|
|
raise TOTP.TOTPException('Invalid secret')
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
args = get_arguments() # Get args from cmdline
|
|
if args.andOTP:
|
|
andOTPfile = args.andOTP
|
|
else:
|
|
andOTPfile = read_config(args) # Read from the cfg file
|
|
data = read_file(andOTPfile) # Open the data file
|
|
|
|
if not args.site: # Show a menu if no input
|
|
while True:
|
|
menu(data)
|
|
site = raw_input("Key #: ")
|
|
try:
|
|
site = int(site)
|
|
except:
|
|
print("Enter a number, not text.")
|
|
|
|
if site > len(data):
|
|
raw_input("Incorrect menu choice, press Enter to try again")
|
|
else:
|
|
print_OTP(data[site-1]["secret"]) # -1 since index start with 0 and the menu with 1
|
|
exit(0)
|
|
|
|
for p in data: # Try to find a matching site
|
|
if p["label"].strip().lower() == args.site[0].lower():
|
|
print_OTP(p["secret"])
|
|
exit(0)
|
|
|
|
|
|
die("Could not find '%s' in the andOTP file" % args.site[0])
|
|
|
|
|