#!/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])