/trackfile/utils.py

https://github.com/9h37/delikatess · Python · 221 lines · 174 code · 12 blank · 35 comment · 5 complexity · 3950a3e9277bf9c6d8d6f0a8ae1cb6a9 MD5 · raw file

  1. # -*- coding: utf-8 -*-
  2. from models import DatabaseEntry
  3. import os
  4. import hashlib
  5. import re
  6. class FileListing:
  7. """ Object use to list files recursively """
  8. def __init__ (self):
  9. pass
  10. def readdir (self, path):
  11. """
  12. List files recursively.
  13. path: Directory to list.
  14. return: List of files containned in @path and its subdirs.
  15. """
  16. ret = []
  17. for filename in os.listdir (path):
  18. abspath = os.path.join (path, filename)
  19. if os.path.isdir (abspath) == True:
  20. ret = ret + self.readdir (abspath)
  21. else:
  22. ret.append (abspath)
  23. return ret
  24. class FileChecksum:
  25. """ Object used to get files' checksum. """
  26. def __init__ (self):
  27. pass
  28. def checksum (self, filelist, block_size = 512):
  29. """
  30. Get checksum of each files
  31. filelist: List of files to hash (can be obtained with the #FileListing object
  32. block_size: Files are read by block of data, this param define the size in bytes of those block (default = 512)
  33. return: A list of tuples (filepath, filehash)
  34. """
  35. ret = []
  36. for f in filelist:
  37. h = hashlib.sha256 ()
  38. try:
  39. fd = open (f, "rb")
  40. except IOError:
  41. print "Can't open :", f
  42. else:
  43. while True:
  44. data = fd.read (block_size)
  45. if not data:
  46. break
  47. h.update (data)
  48. fd.close ()
  49. ret.append ((f, h.hexdigest ()))
  50. return ret
  51. class FileEncrypt:
  52. """ Object used to encrypt a list of files. """
  53. def __init__ (self, recipient, src, dest, nbackups = 5, extra_opts = ""):
  54. """
  55. Object's constructor
  56. recipient: GPG Key used to encrypt files. (need to be generated before the execution)
  57. src: Directory in which files are stored
  58. dest: Directory where to put encrypted files
  59. nbackups: number of copies (default = 5)
  60. """
  61. self.extra_opts = extra_opts
  62. self.recipient = recipient
  63. self.nbackups = nbackups
  64. self.src = src
  65. self.dest = dest
  66. self.pattern = re.compile (r"(\.\d+)*.gpg")
  67. def _backup (self, path, iterate = 1):
  68. """
  69. Internal function which backup existant encrypted files.
  70. path: Path of the original file
  71. iterate: Internal param for the recursivity
  72. """
  73. npath = self.pattern.sub ("." + str (iterate) + ".gpg", path)
  74. if os.path.exists (npath) == True:
  75. self._backup (npath, iterate + 1)
  76. if iterate <= self.nbackups + 1:
  77. # Rename the file
  78. os.rename (path, npath)
  79. # Update the database
  80. try:
  81. dbe = DatabaseEntry.objects.get (path = path)
  82. except DatabaseEntry.DoesNotExist:
  83. pass
  84. else:
  85. dbe.path = npath
  86. dbe.save ()
  87. def encrypt (self, filelist):
  88. """
  89. Encrypt files using GnuPG
  90. filelist: List of files to encrypt (files stored in @self.src)
  91. return: List of tuples (filepath, encryptedpath)
  92. """
  93. ret = []
  94. for f in filelist:
  95. ef = f.replace (self.src, self.dest) + ".gpg"
  96. dirname = os.path.dirname (ef)
  97. if os.path.exists (dirname) == False:
  98. os.makedirs (dirname)
  99. if os.path.exists (ef) == True:
  100. self._backup (ef)
  101. os.system ("gpg " + self.extra_opts + " --output " + ef + " --encrypt --recipient " + self.recipient + " " + f)
  102. ret.append ((f, ef))
  103. return ret
  104. def decrypt (self, filelist):
  105. """
  106. Decrypt files using GnuPG
  107. """
  108. pass
  109. class FileManager:
  110. """ Object implementing all previously created objects. """
  111. def __init__ (self, recipient, src, dest, nbackups = 5, extra_gpg_opts = ""):
  112. """
  113. Object's constructor
  114. See #FileEncrypt constructor for more details.
  115. """
  116. self.src = src
  117. self.fl = FileListing ()
  118. self.fh = FileChecksum ()
  119. self.fe = FileEncrypt (recipient, src, dest, nbackups, extra_gpg_opts)
  120. def _combinedata (self, checksums, encrypted):
  121. """
  122. Combine data in one list of tuples.
  123. checksums: List of tuples: (path, checksum) returned by #FileChecksum.checksum
  124. encrypted: List of tuples: (path, encrypted_path) returned by #FileEncrypt.encrypt
  125. return: List of tuples: (path, checksum, encrypted_path)
  126. """
  127. ret = []
  128. for path,checksum in checksums:
  129. for path2,encrypted_path in encrypted:
  130. if path == path2:
  131. ret.append ((path, checksum, encrypted_path))
  132. break
  133. return ret
  134. def run (self):
  135. """
  136. Do: Listing, hash and encryption.
  137. """
  138. try:
  139. if os.path.exists (self.src) == False:
  140. raise self.src + ": path doesn't exists."
  141. except Exception as e:
  142. print e
  143. else:
  144. files = self.fl.readdir (self.src)
  145. checksums = self.fh.checksum (files)
  146. to_encrypt = files
  147. for dbe in DatabaseEntry.objects.all ():
  148. # Check if file's checksum is already in the database
  149. for f,h in checksums:
  150. # If yes, don't need to encrypt it
  151. if dbe.checksum == h:
  152. to_encrypt.remove (f)
  153. break
  154. # Now, delete unused checksums from the list
  155. for c in checksums:
  156. if c[0] not in to_encrypt:
  157. checksums.remove (c)
  158. # Encrypt only files which are not in the database
  159. encrypted = self.fe.encrypt (to_encrypt)
  160. data = self._combinedata (checksums, encrypted)
  161. # Put all data in the database
  162. for _,filehash,encryptedpath in data:
  163. dbe = DatabaseEntry (checksum = filehash, path = encryptedpath, sent = 0)
  164. dbe.save ()
  165. # vim: tabstop=4 shiftwidth=4 expandtab