import sys
import os
import subprocess
import shutil
import ctypes
import wx
import string
from datetime import datetime
import ctypes

# ensure wxWidgets can be called whenever needed
app = wx.App()
app.MainLoop()

# settings
ourPath = os.path.dirname(os.path.abspath(__file__))
titleStart = 'EncFol - '
robocopyExePath = 'robocopy.exe'
winMergeExePath = 'C:\Program Files (x86)\WinMerge\WinMergeU.exe'
veraCryptPath = 'C:\Program Files\VeraCrypt\VeraCrypt.exe'
veraCryptFormatPath = 'C:\Program Files\VeraCrypt\VeraCrypt Format.exe'
junctionExePath = os.path.join(ourPath, 'Junction64.exe')
bytesPerMB = 1048576
relFreeTooLowSize = 0.06
relFreeTooHighSize = 0.4
growShrinkFactor = 1.3
minSize = 10 * bytesPerMB
minFreeSize = 5 * bytesPerMB

# load the common icons
normalIco = wx.Icon(os.path.join(ourPath, 'Icon.ico'), wx.BITMAP_TYPE_ICO)
attentionIco = wx.Icon(os.path.join(ourPath, 'Attention.ico'), wx.BITMAP_TYPE_ICO)

# gets a backup filepath version to use for the given filepath
def GetBackupFilePath(filePath):
  fileBasePath, fileExtension = os.path.splitext(filePath)
  return fileBasePath + ' - backup ' + datetime.now().strftime('%Y-%m-%d %H-%M-%S') + fileExtension

# shows the given error
def ShowError(title, prompt):
  wx.MessageBox(prompt, titleStart + title, wx.OK | wx.ICON_EXCLAMATION | wx.CENTRE)

# returns if the given drive is mounted
def IsDriveMounted(driveLetter):
  driveNr = ord(driveLetter[0].upper()) - ord('A')
  driveMask = pow(2, driveNr)
  return ctypes.cdll.kernel32.GetLogicalDrives() & driveMask != 0

# gets the next available drive root
# returns None if it couldn't be found
def GetFreeDriveRoot():
  # get a list of all drives' availability state
  driveStatusBitmask = ctypes.cdll.kernel32.GetLogicalDrives()
  driveStatusString = (bin(driveStatusBitmask)[:1:-1] + '0' * 26)[:26]

  # associate drive letters with them
  driveStatusses = zip(string.ascii_uppercase, driveStatusString)

  # weed out the unavailable and only use drive nr. 8 (H:) onward to give USB & DVD drives some space
  availableDriveLetters = ''.join([driveLetter if driveBit == '0' else '' for driveLetter, driveBit in driveStatusses][8:])

  # and return the highest available drive letter
  if len(availableDriveLetters) == 0:
    return None
  else:
    return availableDriveLetters[-1] + ':\\'

# wait dialog to wait for a drive to either be mounted or dismounted
# returns either:
# - ID_OK: drive finished mounting/unmounting
# - ID_CANCEL: user cancelled
class TWaitMountUnmountDlg(wx.Dialog):
  # constructor
  def __init__(self, driveRoot, waitMounted):
    # note the options used
    self.driveRoot = driveRoot
    self.waitMounted = waitMounted

    # init the dialog itself
    wx.Dialog.__init__(self, None, wx.ID_ANY, '{}Waiting for {}'.format(titleStart, 'mounting' if waitMounted else 'unmounting'))
    self.SetIcon(normalIco)

    # set a timer
    self.waitCycleNr = 0
    self.updateTimer = wx.Timer(self)
    self.Bind(wx.EVT_TIMER, self.OnUpdateTimer, self.updateTimer)
    self.updateTimer.Start(500)

    # create the main sizer
    mainBox = wx.BoxSizer(wx.VERTICAL)

    # add the controls
    label = wx.StaticText(self, wx.ID_ANY, label='... Waiting for the encrypted drive to be {} ...'.format('mounted' if waitMounted else 'unmounted'))
    mainBox.Add(label, 0, wx.ALL | wx.EXPAND, 10)
    buttonBox = wx.BoxSizer(wx.HORIZONTAL)
    mainBox.Add(buttonBox, 0, wx.CENTRE, 0)
    cancelCtrl = wx.Button(self, wx.ID_CANCEL, label='Cancel')
    buttonBox.Add(cancelCtrl, 0, 0, 0)

    # and finalize the GUI
    self.SetSizerAndFit(mainBox)
    currentSize = self.GetSize()
    self.SetSizeHints(currentSize.x, currentSize.y, -1, currentSize.y)
    cancelCtrl.SetFocus()
    self.Centre()

  # update timer notifications
  def OnUpdateTimer(self, event):
    # check if the drive is already in the correct state
    if IsDriveMounted(self.driveRoot) == self.waitMounted:
      # yes -> our purpose is fullfilled
      self.EndModal(wx.ID_OK)

# waits for the given drive to be mounted or dismounted
# returns if the drive was mounted or dismounted
def WaitDrive(driveRoot, waitForMount):
  # look if the drive is already in the correct state
  driveStateOK = IsDriveMounted(driveRoot) == waitForMount
  if not driveStateOK:
    # no -> wait for it
    with TWaitMountUnmountDlg(driveRoot, waitForMount) as dlg:
      driveStateOK = wx.ID_OK == dlg.ShowModal()

  # and return the verdict
  return driveStateOK

# inspects if the given container is already mounted
def IsAlreadyMounted(containerMountLogPath, mountPath):
  # check if the mount log is even there
  isMounted = False
  if os.path.isfile(containerMountLogPath):
    # yes -> look at what drive it was mounted last
    file = None
    allReadOK = True
    try:
      file = open(containerMountLogPath, 'r')
      prevDriveRoot = file.readline().strip()
      prevMountPath = file.readline().strip()
    except:
      allReadOK = False
    finally:
      if file != None:
        file.close()
    if allReadOK:
      # done -> look if it's still mounted there
      isMounted = os.path.isdir(prevMountPath) and IsDriveMounted(prevDriveRoot) and mountPath.casefold() == prevMountPath.casefold()

  # and return the verdict
  return isMounted

# asks for the container password to use
def GetContainerPassword():
  passDialog = wx.TextEntryDialog(None, 'Password for Encrypted folder', titleStart + 'Enter the password for the Encrypted folder:', '', wx.OK | wx.CANCEL | wx.TE_PASSWORD | wx.CENTRE)
  if passDialog.ShowModal() == wx.ID_OK:
    return passDialog.GetValue();
  else:
    return ''

# unmounts the given container, possibly also first removing any junction to it
def UnmountContainer(driveRoot, mountPath=None):
  # remove any junction
  allOK = True
  if mountPath != None:
    result = subprocess.run([junctionExePath, '-d', mountPath])
    if result.returncode != 0:
      allOK = False
  if allOK:
    # done -> dismount the container as well
    result = subprocess.run([veraCryptPath, '/dismount', driveRoot[0], '/quit', 'background'])
    if result.returncode != 0:
      allOK = False

  # and return our status
  return allOK

# mounts the given container in the given drive root,
# also creating a temp junction to access it
def MountContainer(containerPath, password, driveRoot, mountPath, containerMountLogPath=None):
  # mount the drive
  allOK = False
  result = subprocess.run([veraCryptPath, '/volume', containerPath, '/letter', driveRoot[0], '/password', password, '/hash', 'sha-512', '/nowaitdlg', '/mountoption', 'timestamp', '/quit', 'background'])
  if result.returncode == 0:
    # done -> wait for it to show up
    if WaitDrive(driveRoot, True):
      # done -> remove any remainders of old junctions
      result = subprocess.run([junctionExePath, '-d', mountPath])

      # and map a junction to it for use under it's own folder name
      result = subprocess.run([junctionExePath, mountPath, driveRoot])
      allOK = result.returncode == 0
      if not allOK:
        # error -> unmount it again
        UnmountContainer(driveRoot)
      else:
        # done -> look if to log the mount result
        if containerMountLogPath != None:
          # yes -> do so
          with open(containerMountLogPath, 'w') as file:
            file.write(driveRoot + '\n' + mountPath)

  # and return if the mount succeeded
  return allOK

# creates the given container with the given size,
# and mounts it in the given drive root,
# also creating a temp junction to access it
def CreateContainer(containerPath, desiredSizeMB, password, driveRoot, mountPath):
  # create the container
  result = subprocess.run([veraCryptFormatPath, '/create', containerPath, '/size', str(desiredSizeMB) + 'M', '/password', password, '/hash', 'sha512', '/encryption', 'AES', '/filesystem', 'NTFS'])

  # and return our status
  return result.returncode == 0

# copies the given folder, also letting the user decide if the copy was successful
def CopyFolder(sourceFolder, destFolder):
  # make the copy
  allOK = False
  result = subprocess.run([robocopyExePath, sourceFolder, destFolder, '/mir', '/e', '/copy:dat', '/dcopy:dat', '/xd', 'System Volume Information', '$RECYCLE.BIN'])

  # let the user decide if the copy is OK
  result = subprocess.run([winMergeExePath, '/r', '/xq', '/u', sourceFolder, destFolder])
  allOK = result.returncode == 0
  if allOK:
    allOK = wx.YES == wx.MessageBox('Is the copied data OK?', titleStart + 'Copied OK', wx.YES_NO | wx.ICON_QUESTION | wx.NO_DEFAULT | wx.CENTRE)

  # and return if all is OK
  return allOK

# single after-mount option
class TAfterMountOption:
  # constructor
  def __init__(self, name):
    self.name = name
    self.commands = []

  # adds the given command
  def AddCommand(self, type, cmdline):
    self.commands.append((type.lower(), cmdline))

# gets the available options from the given options file
# returns a list of commands keyed by their label,
# where command lists are lists of (mode, cmdline)
def GetAllAfterMountOptions(optionsFilePath, mountPath):
  # look if the given options file is present
  allOptions = []
  if not os.path.isfile(optionsFilePath):
    # no -> just return a single option that just opens the mount path
    defaultOption = TAfterMountOption('Default')
    defaultOption.AddCommand('s', mountPath)
    allOptions.append(defaultOption)
  else:
    # yes -> parse it
    currentOption = None
    with open(containerOptionsPath, 'r') as optionsFile:
      for nextLine in optionsFile.readlines():
        nextLine = nextLine.strip()
        if len(nextLine) > 0:
          if nextLine[0] == '[' and nextLine[-1] == ']':
            currentOption = TAfterMountOption(nextLine[1:-1])
            allOptions.append(currentOption)
          elif currentOption != None:
            currentOption.AddCommand(nextLine[0], nextLine[2:].replace("[MOUNTPATH]", mountPath))

  # and return the found options
  return allOptions

# dialog to let the user choose from one of the given options
# returns ID_OK / ID_CANCEL to denote status
class TChooseOptionDlg(wx.Dialog):
  # constructor
  def __init__(self, allOptions):
    # note the options used
    self.allOptions = allOptions
    self.chosenOption = None

    # init the dialog itself
    wx.Dialog.__init__(self, None, -1, '{}What way to mount the folder?'.format(titleStart))
    self.SetIcon(normalIco)

    # create the main sizer
    mainBox = wx.BoxSizer(wx.VERTICAL)

    # add the controls
    label = wx.StaticText(self, wx.ID_ANY, label='What to do?')
    mainBox.Add(label, 0, wx.ALL | wx.EXPAND, 10)
    buttonBox = wx.BoxSizer(wx.VERTICAL)
    mainBox.Add(buttonBox, 0, wx.CENTRE, 0)
    firstOptionCtrl = None
    for nextOption in allOptions:
      optionCtrl = wx.Button(self, wx.ID_ANY, label=nextOption.name)
      if firstOptionCtrl == None:
        firstOptionCtrl = optionCtrl
      optionCtrl.option = nextOption
      optionCtrl.Bind(wx.EVT_BUTTON, self.OnOptionClicked)
      buttonBox.Add(optionCtrl, 0, wx.TOP | wx.CENTRE | wx.EXPAND, 7)
    cancelCtrl = wx.Button(self, wx.ID_CANCEL, label='Cancel')
    buttonBox.Add(cancelCtrl, 0, wx.TOP | wx.CENTRE, 14)

    # and finalize the GUI
    self.SetSizerAndFit(mainBox)
    currentSize = self.GetSize()
    self.SetSizeHints(currentSize.x, currentSize.y, -1, currentSize.y)
    firstOptionCtrl.SetFocus()
    self.Centre()

  # option button click notifications
  def OnOptionClicked(self, event):
    # note the option the button stands for
    self.chosenOption = event.GetEventObject().option

    # and no use to stick around
    self.EndModal(wx.ID_OK)

  # returns the chosen option, or None if none was chosen
  def ChosenOption(self):
    return self.chosenOption

# gets the after-mount option to use
def GetAfterMountOption(containerOptionsPath, mountPath):
  # get all available options
  allOptions = GetAllAfterMountOptions(containerOptionsPath, mountPath)

  # look if there's more than one
  chosenOption = None
  if len(allOptions) == 1:
    # no -> just default to it
    chosenOption = allOptions[0]
  else:
    # yes -> ask for the one to use
    with TChooseOptionDlg(allOptions) as dlg:
      if wx.ID_CANCEL != dlg.ShowModal():
        # done -> use that option
        chosenOption = dlg.ChosenOption()

  # and return the chosen option, if any
  return chosenOption

# dialog to let the user indicate he's done with the container
class TWaitDoneDialog(wx.Dialog):
  # constructor
  def __init__(self, containerName):
    # init the dialog itself
    wx.Dialog.__init__(self, None, wx.ID_ANY, '{}Waiting for {}'.format(titleStart, containerName))
    self.SetIcon(attentionIco)

    # create the main sizer
    mainBox = wx.BoxSizer(wx.VERTICAL)

    # add the controls
    label = wx.StaticText(self, wx.ID_ANY, label="Press 'Done' to close the encrypted folder '{}'".format(containerName))
    mainBox.Add(label, 0, wx.ALL | wx.EXPAND, 10)
    buttonBox = wx.BoxSizer(wx.HORIZONTAL)
    mainBox.Add(buttonBox, 0, wx.CENTER, 0)
    notYetDoneCtrl = wx.Button(self, wx.ID_ANY, label='Not done')
    buttonBox.Add(notYetDoneCtrl, 0, wx.RIGHT, 7)
    doneCtrl = wx.Button(self, wx.ID_ANY, label='Done')
    doneCtrl.Bind(wx.EVT_BUTTON, self.OnDoneClicked)
    buttonBox.Add(doneCtrl, 0, 0, 0)

    # and finalize the GUI
    self.SetSizerAndFit(mainBox)
    currentSize = self.GetSize()
    self.SetSizeHints(currentSize.x, currentSize.y, -1, currentSize.y)
    notYetDoneCtrl.SetFocus()
    self.Centre()

  # done button click notifications
  def OnDoneClicked(self, event):
    self.EndModal(wx.ID_OK)

# runs the commands specified in the given after-mount option
def RunAfterMountCommands(option, mountPath):
  # process each command
  goOn = True
  for nextType, nextCmdLine in option.commands:
    try:
      if nextType == 'r':
        subprocess.run(nextCmdLine, shell=True, stdin=None, stdout=None, stderr=None)
      else:
        os.startfile(nextCmdLine)
    except Exception as exception:
      # error -> tell and quit
      ShowError('Error running command', 'There was an error running the following after-mount command:\n{}.\n\n{}'.format(nextCmdLine, exception))
      goOn = False
    if not goOn:
      break;

# evaluates if the given mounted container is still of the right size
# returns a new suggested size if a resize is needed, or None otherwise
def EvaluateContainerSize(driveRoot, forceResize, containerName):
  # get stats on the container
  oldSize, usedSize, freeSize = shutil.disk_usage(driveRoot)
  relFree = freeSize / oldSize

  # look if the container size is still appropriate
  mustResize = False
  if relFree < relFreeTooLowSize or freeSize < minFreeSize:
    # no (too small) -> note
    mustResize = True
    resizePrePrompt = 'The ' + containerName + ' folder is getting pretty full; expansion is suggested.'
  elif relFree > relFreeTooHighSize:
    # no (too big) -> note
    mustResize = True
    resizePrePrompt = 'The ' + containerName + ' folder is oversized; shrinking is suggested.'
  else:
    # yes -> note
    mustResize = forceResize
    resizePrePrompt = 'Manual resize of ' + containerName + ' is requested.'
  if mustResize:
    # no -> determine the desired new size
    suggestedSize = max(usedSize * growShrinkFactor, minSize)
    suggestedFreeSize = suggestedSize - usedSize
    if suggestedFreeSize < minFreeSize:
      suggestedSize += minFreeSize - suggestedFreeSize
    suggestedSizeMB = round(suggestedSize / bytesPerMB)
    mustResize = suggestedSizeMB != round(oldSize / bytesPerMB) or forceResize
    if mustResize:
      # something else -> ask if allowed to resize
      prompt = resizePrePrompt + '\nSize: {}\nUsed: {}\nFree: {}\n\nResizing the encrypted folder will also empty it\'s recycle bin!\n\nPlease specify the desired size, in MB:'.format(round(oldSize / bytesPerMB), round(usedSize / bytesPerMB), round(freeSize / bytesPerMB))
      dlg = wx.TextEntryDialog(None, prompt, 'Enter new encrypted folder size', str(suggestedSizeMB), wx.OK | wx.CANCEL | wx.CENTRE)
      mustResize = dlg.ShowModal() == wx.ID_OK
      if mustResize:
        desiredSizeMB = dlg.GetValue().strip();
        mustResize = len(desiredSizeMB) > 0 and desiredSizeMB.isdigit()
        if mustResize:
          desiredSizeMB = int(desiredSizeMB)

  # and return the desired size, if any
  return desiredSizeMB if mustResize else None

# gets the invocation details
# returns (operation, containerPath), or (None, None) on error
def GetInvocationDetails():
  # look which container to mount
  allOK = False
  if len(sys.argv) < 3:
    # none -> error
    ShowError('Illegal invocation', 'Wrong invocation...  Need 2 args; <M|R> <encrypted folder path>.')
  else:
    # done -> get the context
    operation = sys.argv[1].lower().strip()[0]
    containerPath = sys.argv[2]

    # look if we're called with a valid operation
    allOK = operation in ['m', 'r']
    if not allOK:
      # no -> error
      ShowError('Illegal invocation', 'Wrong operation specified.  Please only specify either M or R.')

  # and return our result
  if allOK:
    return (operation, containerPath)
  else:
    return (None, None)




# ensure we show the right icon in the taskbar
myappid = 'TwoLogs.MountEncryptedFolder'
ctypes.windll.shell32.SetCurrentProcessExplicitAppUserModelID(myappid)

# look how we're called
operation, containerPath = GetInvocationDetails()
if operation != None:
  # done -> get the container's details
  doFullMount = operation == 'm'
  doResizeOnly = not doFullMount
  mountPath, containerExtension = os.path.splitext(containerPath)
  containerOptionsPath = mountPath + '.encopt'
  containerMountLogPath = mountPath + '.encmnt'
  containerName = os.path.basename(mountPath)

  # look if the container is already mounted
  isAlreadyMounted = IsAlreadyMounted(containerMountLogPath, mountPath)
  if isAlreadyMounted:
    # yes -> look if the operation conflicts
    if doResizeOnly:
      ShowError('Mount conflict', 'This folder is already mounted by a previous invocation; please close that session first.')
      # yes -> tell
    else:
      # no -> look what option to use
      chosenOption = GetAfterMountOption(containerOptionsPath, mountPath)
      if chosenOption != None:
        # done -> run any after-mount commands
        RunAfterMountCommands(chosenOption, mountPath)
  else:
    # no -> look which drive letter is still free
    driveRoot = GetFreeDriveRoot()
    if driveRoot == None:
      # none! -> error
      ShowError('No more drives available', 'All drive letters from H: on are already in use; please close one first before accessing this encrypted folder.')
    else:
      # done -> ask for the password
      password = GetContainerPassword()
      if len(password) > 0:
        # done -> look if to also get any mount options
        if doFullMount:
          # yes -> get one
          chosenOption = GetAfterMountOption(containerOptionsPath, mountPath)

        # look if an option was chosen, when needed
        if doResizeOnly or chosenOption != None:
          # yes -> pre-run Veracrypt so it won't cause subprocess.run to wait on the main veracrypt process
          subprocess.Popen([veraCryptPath, '/q', 'background'])

          # mount the container
          if not MountContainer(containerPath, password, driveRoot, mountPath, containerMountLogPath):
            # it didn't -> error
            ShowError('Encrypted folder not mounted', 'There was an error mounting the corresponding encrypted folder.')
          else:
            # done -> look if to do a full mount
            if doFullMount:
              # yes -> run any after-mount commands
              RunAfterMountCommands(chosenOption, mountPath)

              # and wait for the user to be done with it
              with TWaitDoneDialog(containerName) as dlg:
                dlg.ShowModal()

            # look if a new size is desired
            desiredSizeMB = EvaluateContainerSize(driveRoot, operation == 'r', containerName)
            mustResize = desiredSizeMB != None
            if mustResize:
              # yes -> get a free drive letter to mount the resized container on
              resizeDriveRoot = GetFreeDriveRoot()
              if resizeDriveRoot == None:
                # none! -> error
                ShowError('No more drives available', 'All drive letters from H: on are already in use; please close one first before resizing this encrypted folder.')
                mustResize = False
              else:
                # done -> create the new container, mounting it parallel to the original
                resizeContainerPath = containerPath + '.new'
                resizeMountPath = mountPath + '_new'
                copiedOK = False
                if not CreateContainer(resizeContainerPath, desiredSizeMB, password, resizeDriveRoot, resizeMountPath):
                  # couldn't create -> error
                  ShowError('Error creating new encrypted folder', 'There was an error creating the resized encrypted folder.')
                elif not MountContainer(resizeContainerPath, password, resizeDriveRoot, resizeMountPath):
                  # couldn't mount -> error
                  ShowError('Error mounting the new encrypted folder', 'There was an error mounting the resized encrypted folder.')
                else:
                  # done -> copy everything over
                  copiedOK = CopyFolder(mountPath, resizeMountPath)

                  # and dismount the container again
                  UnmountContainer(resizeDriveRoot, resizeMountPath)

            # remove the junction again
            UnmountContainer(driveRoot, mountPath)
            os.remove(containerMountLogPath)

            # and look if a new container is there
            if mustResize:
              # yes -> wait for the containers to be dismounted
              killResizeCopy = True
              driveFound = WaitDrive(driveRoot, False) and WaitDrive(resizeDriveRoot, False)

              # look if all went well so far
              if driveFound and copiedOK:
                # yes -> backup the old one
                backupFilePath = GetBackupFilePath(containerPath)
                os.rename(containerPath, backupFilePath)

                # and move the new one in it's place
                os.rename(resizeContainerPath, containerPath)
                killResizeCopy = False

              # and kill the resize copy, if still needed
              if killResizeCopy and os.path.isfile(resizeContainerPath):
                # no -> thow away the copy instead
                os.remove(resizeContainerPath)