Initial commit

softtoken/softtoken.py

@@ -0,0 +1,164 @@
+#  Licensed under the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License. You may obtain
+#  a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#  License for the specific language governing permissions and limitations
+#  under the License.
+
+from __future__ import print_function
+
+import argparse
+import base64
+import configparser
+import hashlib
+from os import path
+from os import urandom
+import sys
+
+from pykeyboard import PyKeyboard
+import pyotp
+
+
+__version__ = '0.0.1'
+
+CONFIG_FILE = 'softtoken.conf'
+
+
+def load_config():
+    dir_path = path.join(
+        path.dirname(path.realpath(__file__)),
+        CONFIG_FILE)
+    cfg = configparser.SafeConfigParser()
+    cfg.read(dir_path)
+    return cfg
+
+
+def save_config(cfg):
+    dir_path = path.join(
+        path.dirname(path.realpath(__file__)),
+        CONFIG_FILE)
+    try:
+        with open(dir_path, 'w+') as configfile:
+            cfg.write(configfile)
+    except Exception:
+        print('ERROR: Cannot write config file')
+        sys.exit(2)
+
+
+def create_token(name, hash_function='sha256', digits=6, seed_length=20):
+    cfg = load_config()
+    if cfg.has_section(name):
+        print('Token %s already exists. Delete it first' % name)
+        sys.exit(2)
+
+    seed = urandom(seed_length)
+
+    cfg.add_section(name)
+    cfg.set(name, 'hash', hash_function)
+    cfg.set(name, 'digits', str(digits))
+    cfg.set(name, 'seed', base64.b32encode(seed))
+
+    save_config(cfg)
+
+    print('\nNew Token created:\n\n%s\n-------------' % name)
+    print('Seed (hex): %s' % seed.encode('hex'))
+    print('Seed (b32): %s\n' % base64.b32encode(seed))
+
+
+def delete_token(name):
+    cfg = load_config()
+    if not cfg.has_section(name):
+        print('Token %s does not exist' % name)
+        sys.exit(2)
+    cfg.remove_section(name)
+    save_config(cfg)
+    print("Token %s successfully deleted" % name)
+
+
+def print_tokens():
+    cfg = load_config()
+    for section in cfg.sections():
+        print("[*] %s" % section)
+
+
+def main():
+
+    parser = argparse.ArgumentParser(version=__version__)
+    parser.add_argument('--new', action='store_true', default=False,
+                        dest='new_token', help='Generate a new Soft Token')
+    parser.add_argument('--delete', action='store_true', default=False,
+                        dest='delete_token', help='Delete a Soft Token')
+    parser.add_argument('--list', action='store_true', default=False,
+                        dest='list_tokens', help='List configured tokens')
+    parser.add_argument('--token', '-t', required=False, dest='token_name',
+                        help='Soft Token name')
+    parser.add_argument('--hash', default='sha256', dest='hash_function',
+                        choices=('sha1', 'sha256', 'sha512'), help='Hash '
+                        'function to use (default is sha256)')
+    parser.add_argument('--digits', '-d', type=int, default=6, dest='digits',
+                        help='OTP Length (default is 6)')
+    parser.add_argument('--length', '-l', type=int, default=20,
+                        dest='seed_length', help='Seed length in bytes '
+                        '(default is 20)')
+    parser.add_argument('-X', action='store_true', default=False,
+                        dest='print_focus', help='Output the OTP where '
+                        'the current focus is')
+
+    args = parser.parse_args()
+
+    if args.list_tokens:
+        print_tokens()
+        sys.exit(0)
+
+    if args.token_name is None:
+        print("A Token name is required for this action")
+        parser.print_help()
+        sys.exit(-1)
+
+    if args.new_token:
+        create_token(args.token_name,
+                     args.hash_function,
+                     args.digits,
+                     args.seed_length)
+        sys.exit(0)
+
+    if args.delete_token:
+        delete_token(args.token_name)
+        sys.exit(0)
+
+    if args.list_tokens:
+        print_tokens()
+        sys.exit(0)
+
+    # Generate new OTP if the token exists
+    cfg = load_config()
+    if not cfg.has_section(args.token_name):
+        print('Token %s does not exist' % args.token_name)
+        sys.exit(2)
+
+    if args.hash_function == 'sha1':
+        hf = hashlib.sha256
+    elif args.hash_function == 'sha256':
+        hf = hashlib.sha256
+    elif args.hash_function == 'sha512':
+        hf = hashlib.sha512
+
+    seed = cfg.get(args.token_name, 'seed')
+    totp = pyotp.TOTP(seed, digest=hf, digits=args.digits)
+
+    otp = totp.now()
+
+    if args.print_focus:
+        k = PyKeyboard()
+        k.type_string(otp)
+    else:
+        print(otp)
+
+
+if __name__ == "__main__":
+    main()