#!/usr/bin/python
#
# utility for in-place partition encryption with LUKS container;
# see also:
#
#   http://help.ubuntu.com/community/EncryptedFilesystemHowto7
#
# milki placed this into the PUBLIC DOMAIN.

import os
import sys
import re
from random import randint as rand
from getopt import gnu_getopt as getopt


class luksConvert:

  def run(self):
    print "\033[31m"+"luksConvert v1.0 - NO WARRANTY. USE AT YOUR OWN RISK."+"\033[0;0m\n"
    self.get_args()
    self.testrun_cryptsetup()
    self.check_sdev()
    self.start()

  def get_args(self):

    opts,files = getopt(sys.argv[1:], "yvc:s:h:", ["cipher=","hash=","key-size="])
    opts = dict(opts)
      
    if len(files) != 1:
      print "syntax: luksConvert /dev/unencrypted [...]"
      print "  further options similar to cryptsetup:"
      print "  [-y] -[v] [-c aes-cbc-essiv:sha256] [-h sha1] [-s 128]\n"
      sys.exit()
    else:
      self.sdev = files[0]
      
    self.hash = opts.get("--hash", opts.get("-h", "sha1"))
    self.keysize = opts.get("--key-size", opts.get("-s", "128"))
    self.cipher = opts.get("--cipher", opts.get("-c", "aes-cbc-essiv:sha256"))

    
  def luksCheck(self, devicename):
    output = self.pipe("LC_ALL=C  cryptsetup luksDump " + devicename)
    return output.find("offset:\t1032") >= 0

    
  def testrun_cryptsetup(self):
    print "* performing cryptsetup test run... ",
    fn = "/tmp/lC-test-"+str(rand(1000,9999))
    os.system("dd if=/dev/zero of="+fn+" bs=1000 count=10000 2>/dev/null")
    os.system("mkfs.vfat "+fn+" >/dev/null")
    loop = self.pipe("losetup -f").strip()
    os.system("losetup "+loop+" "+fn+" >/dev/null")
    ok = 0==os.system("echo passw0rd | cryptsetup luksFormat "+loop+" -q -c "+self.cipher+" -h "+self.hash+" -s "+self.keysize+" 2>/dev/null")
    ok = ok and self.luksCheck(loop)
    os.system("losetup -d "+loop+" ; rm "+fn);
    if ok:
      print "ok\n"
    else:
      print "fail\n  WARNING: cryptsetup test run didn't work as expected\n"
      sys.exit()
    return ok


  def check_sdev(self):
    # read
    f = open(self.sdev, "rb")
    startsector = f.read(2048)
    f.close()
    # check magic numbers
    if startsector[0x438:0x43A] != "\x53\xEF":
      print "* "+self.sdev+" doesn't seem to contain ext2/ext3 system\n"
    if startsector[0:4] == "LUKS":
      print "* Source device already has a LUKS header. Don't run this tool twice!"
      print "  (Or put back the backup on, if the last run of luksConvert failed.)\n"
      sys.exit()
    # check filesystem size < below partition size (-1032)
    print "* Hopefully you have shrunken your filesystem on " + self.sdev + "\n  already - by 1032 sectors / 516KB at least! (Use resize2fs.)\n"
    return 0


  def cmd(self, cmd):
    print cmd
    return not os.system(cmd)
    

  def pipe(self, cmd):
    pc = os.popen(cmd, "r")
    pc_s = ""
    try:
      pc_s = pc.read(4000)
    except:
      0
    return pc_s


  # walks through the cryptsetup and sector copying process
  def start(self):

    # read first 5MB
    source = blockdevice(self.sdev, "rb")
    data1 = source.readBlock()

    # back it up
    fn = "luksconvert.backup."+str(rand(1000,9999))
    f = open(fn, "wb")
    f.write(data1)
    f.close()
    print "* First 5MB of source device " + self.sdev + " backed up into\n  "+fn+"\n"
    
    # create encrypted devmapper node
    cryptnode = "lconvert" + str(rand(1000,9999))
    print "* DO NOT ABORT AFTER THIS POINT"
    print "  first sectors will be overwritten, AFTER you entered a password\n"
    if not self.cmd("cryptsetup luksFormat "+self.sdev+" -v -y -c "+self.cipher+" -h "+self.hash+" -s "+self.keysize):
      print "\nERROR. ABORTING.\n"
      sys.exit()

    # check cryptsetup
    if (not self.luksCheck(self.sdev)):
      print "\n* WARNING: the written LUKS header seems to be larger than 1032 sectors. Unexpected. Read buffer (5MB) might have been too small and data already lost. But proceeding anyway. (If you abort here, be sure to put the backup block onto the source device.)\n"
    
    # init devmapper node, open encrypted target
    print "\n* yes, you have to enter your password once more:\n"
    while not self.cmd("cryptsetup luksOpen "+self.sdev+" "+cryptnode):
      print "* ERROR. If you already forgot the password, you need the sectorbackup."
    dest = blockdevice("/dev/mapper/"+cryptnode, "wb")
    
    # copying process
    print "* now copying everything over... please wait\n"
    # 
    # first block is already read,
    # get a second 5MB, before writing the first,
    # because the last 516K data always overlap
    #
    while (1):
      # info
      print "  " + str(source.pos / 1024) + "K\r",
      # pre-read next block
      data2 = source.readBlock()
      if not len(data2):
        break;
      # write
      dest.writeBlock(data1)
      data1 = data2
    # write remaining
    dest.writeBlock(data1)
    print "            \n"
    
    # close up
    dest.handle.close()
    source.handle.close()
    self.cmd("cryptsetup luksClose "+cryptnode)
    print
    self.cmd("sync")

    # thanks for buying
    print "\n* done        \n"
      

class blockdevice:

  def __init__(self, fn, mode, bsize=5*1024*1024):
    self.fn = fn
    self.handle = open(fn, mode)
    self.bsize = bsize
    self.pos = 0
    
  def readBlock(self):
    try:
      block = self.handle.read(self.bsize)
    except:
      print "eof / read error"
    self.pos += len(block)
    return block
    
  def writeBlock(self, block):
    try:
      self.handle.write(block)
    except Exception, e:
      print
      print e
      print "\n* writing error or device "+self.fn+" full"
      print "  (should in fact occour for the last 1032 sectors from source device)"
    self.pos += len(block)

     


luksConvert().run()