#!/usr/bin/env python3

import argparse
import getpass
import keyring
import os
import subprocess
import sys

DESCRIPTION = "Wrapper for VeraCrypt to assist in working with data on an encrypted usb drive (e.g.\ a GIT repository)."
KEYRING_SERVICE = "de.tu-dortmund.cs.ls1.veracryptwrapper.hisdnp"
VERA_CONTAINER_NAME_ENV_VAR = "VERA_CONTAINER_NAME"

DEFAULT_DRIVE_NAME = "KINGSTON"
DEFAULT_VERA_VOL_NAME = "NO NAME"
#TODO determine this based on OS
DEFAULT_MNT_PATH = "/Volumes"

VERA_BINARIES = [
        'veracrypt',
        'VeraCrypt',
        '/Applications/VeraCrypt.app/Contents/MacOS/VeraCrypt', #MacOS
]

VERA_PARAM_TEXT_MODE = '-t'
VERA_PARAM_STDIN = '--stdin'
VERA_PARAM_NON_INTERACTIVE = '--non-interactive'

ARG_CONTAINER_NAME = "arg_container"
ARG_USB_DRIVE_NAME = "arg_usbstick"
ARG_VERA_VOL_NAME = "arg_volname"
ARG_MOUNT_PATH = "arg_mountpath"
ARG_UNMOUNT_USB_DRIVE = "arg_unmountusbdrive"
ARG_RESET_PASSWORD = "arg_resetpw"
ARG_VERA_DISMOUNT = "arg_vera_dismount"

def __vera_binary():
    binaries = [ b for b in VERA_BINARIES if os.path.exists(b) ]
    assert binaries, "Unable to find VeraCrypt executable!"

    return binaries[0]

def __sync_io():
    #Sync all I/O devices (just in case)
    res = subprocess.run([
        'sync'
    ])
    res.check_returncode()
    
def __store_password(username, pw):
    keyring.set_password(KEYRING_SERVICE, username, pw)

def __delete_password(username):
    keyring.delete_password(KEYRING_SERVICE, username)

def __get_password(username, reset_pw = False):
    if reset_pw:
        __delete_password(username)

    pw = keyring.get_password(KEYRING_SERVICE, username)
    if not pw:
        pw = getpass.getpass()
    return pw

def __vera_exec(option_list, password = None):
    args = [ __vera_binary(), VERA_PARAM_TEXT_MODE ]

    if password:
        args = args + [ VERA_PARAM_STDIN ]
        password = password.encode()

    return subprocess.run(args + option_list, input = password).returncode

def vera_unmount(volume_path):
    __sync_io()
    res = __vera_exec(['-d', volume_path])
    assert res == 0, \
            "Failed to dismount VeraCrypt volume! VeraCrypt exited with status {}.".format(res)
    __sync_io()

#TODO the following function is MacOS specific
def unmount_usbdrive(usbdrive_mnt_path):
    __sync_io()
    res = subprocess.run([
        'diskutil',
        'unmount',
        usbdrive_mnt_path,
    ])

    res.check_returncode()

def vera_mount(volume_path, container, identifier):
    pw = __get_password(identifier)

    res = __vera_exec([
        VERA_PARAM_NON_INTERACTIVE,
        '--mount', container,
    ], pw)

    assert res == 0, \
            "Failed to mount VeraCrypt volume! VeraCrypt exited with status {}.".format(res)
    __store_password(identifier, pw)

def __parse_cmdline():
    parser = argparse.ArgumentParser(description = DESCRIPTION)

    container_name_default = os.environ.get(VERA_CONTAINER_NAME_ENV_VAR)
    if container_name_default:
        parser.add_argument(ARG_CONTAINER_NAME, default = container_name_default, nargs = '?',
            help = "Path of the encrypted container file on the usb drive")
    else:
        parser.add_argument(ARG_CONTAINER_NAME,
            help = "Path of the encrypted container file on the usb drive")

    parser.add_argument('--mount-path', '-m', dest = ARG_MOUNT_PATH, default = DEFAULT_MNT_PATH)
    parser.add_argument('--vera-volume-name', '-v', dest = ARG_VERA_VOL_NAME, default = DEFAULT_VERA_VOL_NAME)
    parser.add_argument('--usb-drive-part-name', '-n', dest = ARG_USB_DRIVE_NAME, default = DEFAULT_DRIVE_NAME)

    parser.add_argument('--reset-password', '-p', dest = ARG_RESET_PASSWORD, action = 'store_true')
    parser.add_argument('--dismount', '-d', dest = ARG_VERA_DISMOUNT, action = 'store_true')
    parser.add_argument('--unmount-usb-drive', '-u', dest = ARG_UNMOUNT_USB_DRIVE, action = 'store_true')

    return vars(parser.parse_args())

def __main():
    args = __parse_cmdline()

    volume_name = args[ARG_VERA_VOL_NAME]
    mnt_path = args[ARG_MOUNT_PATH]
    volume = os.path.join(mnt_path, volume_name)

    usbdrive_name = args[ARG_USB_DRIVE_NAME]
    usbdrive = os.path.join(mnt_path, usbdrive_name)

    container_name = args[ARG_CONTAINER_NAME]
    container = os.path.join(usbdrive, container_name)

    identifier = ':'.join([
        getpass.getuser(),
        usbdrive_name,
        container_name,
    ])
    
    if args[ARG_RESET_PASSWORD]:
        __delete_password(identifier)

    if not os.path.exists(usbdrive):
        print(":: USB partition mount point \"{}\" doesn't exist. Ignoring request.".format(usbdrive))
        return

    if args[ARG_VERA_DISMOUNT]:
        if os.path.exists(volume):
            vera_unmount(volume)
        else:
            print(":: VeraCrypt mount point \"{}\" doesn't exist. Skipping unmount!".format(volume))

        if args[ARG_UNMOUNT_USB_DRIVE]:
            unmount_usbdrive(usbdrive)
        return

    if os.path.exists(volume):
        print(":: VeraCrypt mount point \"{}\" already exists! Assuming container is already mounted.".format(volume))
        return

    vera_mount(volume, container, identifier)

if __name__ == '__main__':
    __main()