#!/usr/bin/env python
import argparse
import configparser
import os
import pyotp


def create_config(account_name):
    config = configparser.ConfigParser()
    config.add_section(account_name)
    otp_token = raw_input('Enter OTP Token: ')
    config.set(account_name, 'otp_token', otp_token)
    mode = 'ab'
    if not os.path.exists(os.path.expanduser('~/.otp.conf')):
      mode = 'wb'
    with open(os.path.expanduser('~/.otp.conf'), mode) as configfile:
      config.write(configfile)
    return True


def onetimepass(token):
    return pyotp.TOTP(token).now()


def get_otp(config_path='~/.otp.conf', service=None):
    config = configparser.ConfigParser()
    config.read([os.path.expanduser(config_path)])
    token = config.get(service, 'otp_token')
    otp = onetimepass(token)
    return otp

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='OTP Generator'
    )

    parser.add_argument(
        '--add', '-a',
        required=False,
        #action='store_true',
        help='--add -a: Add otp account',
        dest='add',
    )

    parser.add_argument(
        '--service', '-s',
        required=False,
        help='--service -s: Service',
        dest='service',
    )
    args = vars(parser.parse_args())
    if args['add']:
        print(create_config(account_name=args['add']))
    else:
        service = args['service']
        if not service:
            print('No service specified. Exiting.')
            exit(1)
        print(get_otp(service=args['service']))