diskIDMain

#!/usr/bin/env python2
from PyQt4 import QtGui, QtCore
from distutils.dir_util import copy_tree
import sys
import os
import diskIDMainGUI
import chooseFormats 
import subprocess
import json
import shutil
import sys
import getpass
from subprocess import Popen
from subprocess import check_output
from itertools import islice


global inpDirLoc
global outDirLoc
global dtcDirloc


impDirLoc = ""
outDirLoc = ""
dtcDirLoc = ""

spath = os.getenv("HOME") + '/.diskFormatID/settings.json'

class kryoMain(QtGui.QMainWindow, diskIDMainGUI.Ui_kryoMain):
  def __init__(self, parent=None):
    #don't really understand global variables well enough, but I think I need to put these here
    global inpDirLoc
    
    global outDirLoc
    
    global dtcDirLoc
    # start the whole gui thing
    super(kryoMain, self).__init__(parent)
    self.setupUi(self)
    
    # set the dtc max concurrent isntances box to default to 1
    self.DTCInstances.setText("1")

    # Check for existence of settings file for local user. If it doesn't
    # exist (e.g. on first use) warn and create)

    if not os.path.exists(spath):

      QtGui.QMessageBox.warning(self, "First run notice",
                    "It looks like you're running the program for the first time as this user. I'm creating a new default settings file for you in " + spath,
                    QtGui.QMessageBox.Ok, QtGui.QMessageBox.NoButton,
                    QtGui.QMessageBox.NoButton)

      os.mkdir(os.getenv("HOME") + '/.diskFormatID')
      sfile = open(spath, 'w+')
      sfile.write('{"outputFormats": [{"Image type": "-i4"}, {"Image type": "-i10"}, {"Image type": "-i9"}], "mainSettings": [{"Output directory": "/home/euan/kryo/output/", "DTC directory": "/home/euan/kryo/kryoflux_2.6_linux/dtc/", "Input directory": "/home/euan/kryo/smallTest/", "Delete unmountable": "TRUE"}]}')
      sfile.close()

    # Probably need to change ownership settings if this program 
    # is being run as root, and the user later decides to run it 
    # as themselves. Fix this ina  future version. Right now, 
    # the settings directory and file installed in the dot-directory 
    # will be root-owned if run as sudo.

    #open the settings.json file for reading
    if os.stat(spath).st_size != 0:
      with open(spath) as settings_file:    
      
    #set main settings to values from settings json file
        settings = json.load(settings_file)
        if "mainSettings" in settings:
          for setting in settings["mainSettings"]:
            self.inputDir.setText(setting["Input directory"])
            dtcDirLoc = str(setting["DTC directory"])
            self.DTCDir.setText(setting["DTC directory"])
            inpDirLoc = str(setting["Input directory"])
            self.outputDir.setText(setting["Output directory"])
            outDirLoc = str(setting["Output directory"])
            if setting == "Delete unmountable":
              if setting["Delete unmountable"] == "TRUE":
                self.deleteUnmountable.setChecked(true)
              else:
                  self.deleteUnmountable.setChecked(true)
                 


# initiate sub program to choose formats when button is clicked      
       
  def choose_formats(self):
    
    self.addTypes = chooseFormats.addTypes()
    self.addTypes.show()

  
# function below takes input from create images button and creates the images

  
  def create_images(self):
    global inpDirLoc
    
    global outDirLoc
    
    global dtcDirLoc
    
    with open(spath) as data_file:    
      data = json.load(data_file)
       
    # set input variables for the source and destination file directories and dtc directory
    
    source_dirname= inpDirLoc
    dest_dir = outDirLoc
    dtc_dir = dtcDirLoc
    commands = []
    
    outdir = ()
    
    # check the input folder and find all sub directories
    for dirname, dirnames, filenames in os.walk(source_dirname):
    
    # get the folder names for each sub directory
      for subdirname in dirnames:
    
    # create a variable to use for the sub directory paths
        indirname = os.path.join(dirname, subdirname)
        #print str(indirname)
      
   
    # check each sub directory to see if it has a stream file in it
        if os.path.isfile(os.path.join(indirname, "track00.0.raw")) == True:
    
    # create directories to copy the streamfiles to, as the directories have to have no spaces in them, then copy the files to those directories
   
          outTrackDir = os.path.join(dest_dir,subdirname.replace(" ","_"),"track")
          if not os.path.exists(outTrackDir):
            os.makedirs(outTrackDir)
          outdir = os.path.join(dest_dir, subdirname.replace(" ","_"))

          for trackFile in os.listdir(indirname):
            if not os.path.isfile(os.path.join(outTrackDir,trackFile)):
              self.imageCreationResults.append("copying track " + str(os.path.join(indirname,trackFile)))
              QtGui.QApplication.processEvents()
              shutil.copyfile(os.path.join(indirname,trackFile),os.path.join(outTrackDir,trackFile))
 
    
    
          
  #setup a variable to form each command into    
          for format in data["outputFormats"]:
            command = set()
            
            filepath = os.path.join(outdir, subdirname.replace(" ","_"))
            #print str(trackPath)
      
            for param in format:
              if command != "":        
                command.add(format[param])
                filepath = filepath + format[param]
                                    
              else:
                command = format[param]   
                filepath = filepath + format[param]            
        
            commands.append((dtc_dir + "./dtc", "-f" + str(filepath) + ".img", "-m1", ", ".join([str(x) for x in command]), "-f" + outTrackDir + "/track","-i0"))
    

    			
    # Now run the commands in paralell
    # set max concurrent instances from feild in gui
    max_workers = int(self.DTCInstances.text())  
    processes = (Popen(cmd, cwd=dtc_dir,stdout=subprocess.PIPE,stderr=subprocess.PIPE) for cmd in commands)

    running_processes = list(islice(processes, max_workers))  
    while running_processes:
      results = []
      for i, process in enumerate(running_processes):
        out, err = process.communicate()
        
#update the text browser with the status
        
        
        results.append(str(out) + str(err))
# process the update to the gui
        QtGui.QApplication.processEvents()
        if process.poll() is not None:  # the process has finished
          running_processes[i] = next(processes, None)  # start new process
          if running_processes[i] is None: # no new processes
            del running_processes[i]
            break
      for result in results:
        self.imageCreationResults.append(result)
    self.imageCreationResults.append(" ")
    self.imageCreationResults.append("------------------------------------------------")
    self.imageCreationResults.append("Creation of images complete")
    self.imageCreationResults.append("------------------------------------------------")
    

  def try_mounting(self):
    global outDirLoc
    dest_dir = outDirLoc
   
        # now try mounting the files you created and delete the ones that don't mount
    
    #select all directories in the destination directory
    for dirname, dirnames, filenames in os.walk(dest_dir):
    
    # for each of those find the files in them
      for subdirname in dirnames:
        indirname = os.path.join(dirname,subdirname)
        
    
    #for each disk image file in that directory
        for dirname2, dirnames2, filenames2 in os.walk(indirname):
          for filenm in filenames2:
            if not "track" in str(os.path.join(dest_dir, indirname, filenm)):
    # for each file, mount it and put the resulting errors into a variable p.stderr
     #update the text browser with the status
              self.MountingResults.append("Trying to mount " + str(filenm))
# process the update to the gui
              QtGui.QApplication.processEvents()
              p = subprocess.Popen(["mount", "-O loop", os.path.join(dest_dir, indirname, filenm), "/mnt"], stderr=subprocess.PIPE)
    
    # if there is no error then umount the disk image
              if p.stderr.readline() == "":
     #update the text browser with the status
                self.MountingResults.append(str(filenm) + " mounted successfully")
# process the update to the gui
                QtGui.QApplication.processEvents()
                subprocess.call(["umount", "/mnt"])
                self.MountingResults.append(" ")
                self.MountingResults.append("------------------------------------------------")
                self.MountingResults.append(" ")
  
                   
              else:
#update the text browser with the status
                self.MountingResults.append(str(filenm) + " mounting failed")

# process the update to the gui
                QtGui.QApplication.processEvents()
                if self.deleteUnmountable.isChecked():
#update the text browser with the status
                  self.MountingResults.append("Deleting " + str(filenm))
                  self.MountingResults.append(" ")
                  self.MountingResults.append("------------------------------------------------")
                  self.MountingResults.append(" ")
# process the update to the gui
                  QtGui.QApplication.processEvents()
                  subprocess.call(["rm", os.path.join(dest_dir, indirname, filenm)])
                else:
                
                  self.MountingResults.append(" ")
                  self.MountingResults.append("------------------------------------------------")
                  self.MountingResults.append(" ")
    self.MountingResults.append
    self.MountingResults.append("Mounting process complete")
    self.MountingResults.append(" ")
    self.MountingResults.append("------------------------------------------------")       
  
  def delete_unmountable_change(self):
    
    if os.stat(spath).st_size != 0:
# if it is not empty open it for reading
      with open(spath,'r') as settings_file: 
        settings = json.load(settings_file)
        if self.deleteUnmountable.isChecked():
          for setting in settings["mainSettings"]:
            setting["Delete unmountable"] = "TRUE"
          
        else: 
          for setting in settings["mainSettings"]:
            setting["Delete unmountable"] = "FALSE"
        jsoutput = json.dumps(settings)
        outfile = open(spath, 'w')
        outfile.write(jsoutput)
        outfile.close()

  def browse_dtc_dir(self):
    
    global dtcDirLoc
    dtcDirLoc = QtGui.QFileDialog.getExistingDirectory(self,
                                                           "Select a folder") + "/"
    dtcDirLoc = str(dtcDirLoc)
    self.DTCDir.setText(dtcDirLoc)
    
    if os.stat(spath).st_size != 0:
# if it is not empty open it for reading
      with open(spath,'r') as settings_file: 
        settings = json.load(settings_file)
        for setting in settings["mainSettings"]:
           setting["DTC directory"] = str(dtcDirLoc)

        jsoutput = json.dumps(settings)
        outfile = open(spath, 'w')
        outfile.write(jsoutput)
        outfile.close()

#This function processes the browse input directory button press
  def browse_inp_dir(self):
    global inpDirLoc
    
    inpDirLoc = QtGui.QFileDialog.getExistingDirectory(self,
                                                           "Select a folder") + "/"
    inpDirLoc = str(inpDirLoc)
    self.inputDir.setText(inpDirLoc)
    
    if os.stat(spath).st_size != 0:
# if it is not empty open it for reading
      with open(spath,'r') as settings_file: 
        settings = json.load(settings_file)
        for setting in settings["mainSettings"]:
           setting["Input directory"] = str(inpDirLoc)

        jsoutput = json.dumps(settings)
        outfile = open(spath, 'w')
        outfile.write(jsoutput)
        outfile.close()
  
  def save_mounting_results(self):
    # get file path from user input browsing
    saveFilePath = QtGui.QFileDialog.getSaveFileName(self,'Choose save file location',str(outDirLoc),selectedFilter='*.txt')
    #make file extension .txt if not already
    if str(saveFilePath)[-4:] != '.txt':
      saveFilePath = saveFilePath + ".txt"
    #assing textbrowser contents to a variable and write it to the file selected above
    mountingResultsDialog = self.MountingResults.toPlainText()
    outfile = open(saveFilePath, 'w')
    outfile.write(mountingResultsDialog)
    outfile.close()
   



  def save_dtc_results(self):
     # get file path from user input browsing
    saveDTCPath = QtGui.QFileDialog.getSaveFileName(self,'Choose save file location',str(outDirLoc),selectedFilter='*.txt')
    #make file extension .txt if not already
    if str(saveDTCPath)[-4:] != '.txt':
      saveDTCPath = saveDTCPath + ".txt"
    #assing textbrowser contents to a variable and write it to the file selected above
    imageCreationResultsDialog = self.imageCreationResults.toPlainText()
    outfile = open(saveDTCPath, 'w')
    outfile.write(imageCreationResultsDialog)
    outfile.close()

  
  def browse_output_dir(self):
    global outDirLoc
    
    outDirLoc = QtGui.QFileDialog.getExistingDirectory(self,
                                                           "Select a folder") + "/"
    outDirLoc = str(outDirLoc)
    self.outputDir.setText(outDirLoc)
    
    if os.stat(spath).st_size != 0:
# if it is not empty open it for reading
      with open(spath,'r') as settings_file: 
        settings = json.load(settings_file)
        for setting in settings["mainSettings"]:
           setting["Output directory"] = str(outDirLoc)
         
        jsoutput = json.dumps(settings)
        outfile = open(spath, 'w')
        outfile.write(jsoutput)
        outfile.close()

def main():
    app = QtGui.QApplication(sys.argv)
    form = kryoMain()
    form.show()
    app.exec_()
  

	

if __name__ == '__main__':
    main()