#!/usr/bin/env python
"""
# Simple password manager - "PwManager" import
# Copyright (c) 2011 Michael Buesch <m@bues.ch>
# Licensed under the GNU/GPL version 2 or later.
"""

import sys
import getopt
import libpwman as pwman
import hashlib
import Crypto.Cipher.Blowfish as Blowfish
import zlib
import bz2
import xml.parsers.expat as expat
from base64 import b64decode


opt_dryRun = False
opt_verbose = False
opt_forceOverwrite = False


class Error(Exception): pass

class Entry(object):
	def __init__(self, category):
		self.category = category
		self.desc = None
		self.name = None
		self.pw = None
		self.comment = None
		self.url = None
		self.launcher = None
		self.lvp = None
		self.binary = None
		self.base64 = False

	def interpretData(self):
		if self.desc is None or self.name is None or\
		   self.pw is None or self.comment is None or\
		   self.url is None or self.launcher is None or\
		   self.lvp is None or self.binary is None:
			raise Error("XML data did not contain all fields")
		if not self.desc.strip():
			self.desc = ""
		if not self.name.strip():
			self.name = ""
		if not self.pw.strip():
			self.pw = ""
		if not self.comment.strip():
			self.comment = ""
		if not self.url.strip():
			self.url = ""
		if not self.launcher.strip():
			self.launcher = ""
		if self.base64:
			try:
				self.desc = b64decode(self.desc)
				self.name = b64decode(self.name)
				self.pw = b64decode(self.pw)
				self.comment = b64decode(self.comment)
				self.url = b64decode(self.url)
				self.launcher = b64decode(self.launcher)
			except (TypeError), e:
				raise Error("Base64 encoding error: " + str(e))

class Context(object):
	def __init__(self):
		self.lvl = LVL_ROOT
		self.curCatName = None
		self.curEntry = None
		self.entries = []

LVL_ROOT	= 0
LVL_P		= 1
LVL_CATLIST	= 2
LVL_CAT		= 3
LVL_ENT		= 4
LVL_DESC	= 5
LVL_NAME	= 6
LVL_PW		= 7
LVL_COMM	= 8
LVL_URL		= 9
LVL_LAU		= 10
LVL_LVP		= 11
LVL_BIN		= 12

def pwmImportEntry(p, entry):
	global opt_forceOverwrite

	entry.interpretData()
	bulk = []
	if entry.comment:
		bulk.append(entry.comment)
	if entry.url:
		bulk.append("URL: " + entry.url)
	if entry.launcher:
		bulk.append("Program: " + entry.launcher)
	if entry.binary.strip() == "1":
		bulk.append("Entry contains binary data")
	bulk = "\n".join(bulk)
	e = pwman.PWManEntry(category=entry.category,
			     title=entry.desc,
			     user=entry.name,
			     pw=entry.pw,
			     bulk=bulk)
	if p.entryExists(e):
		if not opt_forceOverwrite:
			print "\n\nThis entry does already exist in the database:"
			print e.dump("\t")
			while True:
				inp = raw_input("Overwrite? [y/n/a] ")
				if inp.lower() == "n":
					print "Skipping."
					return
				if inp.lower() == "y":
					break
				if inp.lower() == "a":
					opt_forceOverwrite = True
					break
		p.delEntry(e)
	p.addEntry(e)
	if opt_verbose:
		print "Imported entry:"
		print e.dump("\t") + "\n"

def pwmParseXml(p, xmlData):
	c = Context()

	def unkElem(name):
		raise Error("Unknown XML element " + name)
	def unkEnd(name):
		raise Error("Unknown XML-end element " + name)
	def unkData(data):
		raise Error("Unexpected XML data " + data)

	def startElement(name, attrs):
		try:
			if c.lvl == LVL_ROOT:
				if name == "P":
					c.lvl = LVL_P
				else:
					unkElem(name)
			elif c.lvl == LVL_P:
				if name == "c":
					c.lvl = LVL_CATLIST
				else:
					unkElem(name)
			elif c.lvl == LVL_CATLIST:
				if name.startswith("c") and name[1:].isdigit():
					c.lvl = LVL_CAT
					c.curCatName = attrs["n"]
				else:
					unkElem(name)
			elif c.lvl == LVL_CAT:
				if name.startswith("e") and name[1:].isdigit():
					c.lvl = LVL_ENT
					c.curEntry = Entry(c.curCatName)
					if "b64" in attrs.keys() and\
					   attrs["b64"].strip() == "1":
						c.curEntry.base64 = True
				else:
					unkElem(name)
			elif c.lvl == LVL_ENT:
				if name == "d":
					c.lvl = LVL_DESC
				elif name == "n":
					c.lvl = LVL_NAME
				elif name == "p":
					c.lvl = LVL_PW
				elif name == "c":
					c.lvl = LVL_COMM
				elif name == "u":
					c.lvl = LVL_URL
				elif name == "l":
					c.lvl = LVL_LAU
				elif name == "v":
					c.lvl = LVL_LVP
				elif name == "b":
					c.lvl = LVL_BIN
				else:
					unkElem(name)
			else:
				assert(0)
		except (KeyError, AttributeError), e:
			raise Error("XML error at %s / %s" % (name, str(attrs)))

	def endElement(name):
		if c.lvl == LVL_ROOT:
			unkEnd(name)
		elif c.lvl == LVL_P:
			if name == "P":
				c.lvl = LVL_ROOT
			else:
				unkEnd(name)
		elif c.lvl == LVL_CATLIST:
			if name == "c":
				c.lvl = LVL_P
			else:
				unkEnd(name)
		elif c.lvl == LVL_CAT:
			if name.startswith("c") and name[1:].isdigit():
				c.lvl = LVL_CATLIST
				c.curCatName = None
			else:
				unkEnd(name)
		elif c.lvl == LVL_ENT:
			if name.startswith("e") and name[1:].isdigit():
				c.lvl = LVL_CAT
				c.entries.append(c.curEntry)
				c.curEntry = None
			else:
				unkEnd(name)
		elif c.lvl == LVL_DESC:
			if name == "d":
				c.lvl = LVL_ENT
			else:
				unkEnd(name)
		elif c.lvl == LVL_NAME:
			if name == "n":
				c.lvl = LVL_ENT
			else:
				unkEnd(name)
		elif c.lvl == LVL_PW:
			if name == "p":
				c.lvl = LVL_ENT
			else:
				unkEnd(name)
		elif c.lvl == LVL_COMM:
			if name == "c":
				c.lvl = LVL_ENT
			else:
				unkEnd(name)
		elif c.lvl == LVL_URL:
			if name == "u":
				c.lvl = LVL_ENT
			else:
				unkEnd(name)
		elif c.lvl == LVL_LAU:
			if name == "l":
				c.lvl = LVL_ENT
			else:
				unkEnd(name)
		elif c.lvl == LVL_LVP:
			if name == "v":
				c.lvl = LVL_ENT
			else:
				unkEnd(name)
		elif c.lvl == LVL_BIN:
			if name == "b":
				c.lvl = LVL_ENT
			else:
				unkEnd(name)
		else:
			assert(0)

	def charData(data):
		if c.lvl == LVL_ROOT:
			unkData(data)
		elif c.lvl == LVL_P:
			unkData(data)
		elif c.lvl == LVL_CATLIST:
			unkData(data)
		elif c.lvl == LVL_CAT:
			unkData(data)
		elif c.lvl == LVL_ENT:
			unkData(data)
		elif c.lvl == LVL_DESC:
			c.curEntry.desc = data
		elif c.lvl == LVL_NAME:
			c.curEntry.name = data
		elif c.lvl == LVL_PW:
			c.curEntry.pw = data
		elif c.lvl == LVL_COMM:
			c.curEntry.comment = data
		elif c.lvl == LVL_URL:
			c.curEntry.url = data
		elif c.lvl == LVL_LAU:
			c.curEntry.launcher = data
		elif c.lvl == LVL_LVP:
			c.curEntry.lvp = data
		elif c.lvl == LVL_BIN:
			c.curEntry.binary = data
		else:
			assert(0)

	parser = expat.ParserCreate("UTF-8")
	parser.StartElementHandler = startElement
	parser.EndElementHandler = endElement
	parser.CharacterDataHandler = charData
	parser.Parse(xmlData, True)

	for entry in c.entries:
		pwmImportEntry(p, entry)
	if opt_dryRun:
		p.flunkDirty()
	else:
		p.commit()

def pwmBlowfishDecrypt(data, pw):
	if len(data) % 8:
		raise Error("Invalid payload length")
	i = 0
	dec = []
	# Bug: Re-init cipher for every block
	while i < len(data):
		fb = Blowfish.new(pw, Blowfish.MODE_CBC)
		dec.append(fb.decrypt(data[i:i+8]))
		i += 8
	data = "".join(dec)
	# unpad
	i = len(data) - 1
	while i and data[i] != '\x01':
		i -= 1
	data = data[0:i]
	return data

def pwmImport(p, pwmFile):
	try:
		pwmData = file(pwmFile, "rb").read()
	except (IOError), e:
		raise Error("Failed to read file '%s': %s" %\
			(pwmFile, e.strerror))
	try:
		if pwmData[0:17] != "PWM_PASSWORD_FILE":
			raise Error("%s: Invalid file header." % pwmFile)
		if ord(pwmData[17]) != 0x05:
			raise Error("%s: Unsupported file version." % pwmFile)
		if ord(pwmData[18]) != 0x01:
			raise Error("%s: Unsupported key hash algorithm." % pwmFile)
		if ord(pwmData[19]) != 0x01:
			raise Error("%s: Unsupported data hash algorithm." % pwmFile)
		if ord(pwmData[20]) != 0x01:
			raise Error("%s: Unsupported crypto algorithm." % pwmFile)
		if ord(pwmData[21]) not in (0x00, 0x01, 0x02):
			raise Error("%s: Unsupported compression algorithm." % pwmFile)
		if ord(pwmData[22]) != 0x00:
			raise Error("%s: Unsupported master password type "
				"(keycard not supported)." % pwmFile)
		keyHash = pwmData[87:107]
		dataHash = pwmData[107:127]
		data = pwmData[127:]
	except (IndexError), e:
		raise Error("%s file format error" % pwmFile)
	pw = pwman.readPassphrase("%s master password" % pwmFile)
	if not pw:
		raise Error("No master password given")
	if keyHash != hashlib.sha1(pw).digest():
		raise Error("%s: Wrong master password" % pwmFile)
	data = pwmBlowfishDecrypt(data, pw)
	if ord(pwmData[21]) == 0x01:
		data = zlib.decompress(data)
	elif ord(pwmData[21]) == 0x02:
		data = bz2.decompress(data)
	if dataHash != hashlib.sha1(data).digest():
		raise Error("%s: Corrupt data" % pwmFile)
	pwmParseXml(p, data)

def usage():
	print "pwman v%d - 'PwManager' import" % pwman.VERSION
	print ""
	print "Usage: %s [OPTIONS] file.pwm" % sys.argv[0]
	print ""
	print "Options:"
	print " -d|--database PATH      Use PATH as database file."
	print "                         If not given, %s is used." %\
						pwman.getDefaultDatabase()
	print " -N|--dry-run            Don't write any data to the database."
	print " -V|--verbose            Print verbose messages."

def main():
	global opt_dryRun
	global opt_verbose

	dbFile = None
	try:
		(opts, args) = getopt.getopt(sys.argv[1:],
			"hd:NV",
			[ "help", "database=", "dry-run", "verbose", ])
		for (o, v) in opts:
			if o in ("-h", "--help"):
				usage()
				return 0
			if o in ("-d", "--database"):
				dbFile = v
			if o in ("-N", "--dry-run"):
				opt_dryRun = True
			if o in ("-V", "--verbose"):
				opt_verbose = True
	except (getopt.GetoptError), e:
		usage()
		return 1
	pwmFiles = args
	if not pwmFiles:
		print "No .pwm file specified"
		return 1
	if not dbFile:
		dbFile = pwman.getDefaultDatabase()
	if not dbFile:
		print "No database file specified"
		return 1
	print "Opening database '%s'..." % dbFile
	passphrase = pwman.readPassphrase("Master passphrase",
					  not pwman.fileExists(dbFile))
	if passphrase is None:
		return 1
	try:
		p = pwman.PWMan(dbFile, passphrase)
		for pwmFile in pwmFiles:
			print "Importing '%s' into database '%s'..." % (pwmFile, dbFile)
			pwmImport(p, pwmFile)
	except (pwman.PWMan.Error, Error), e:
		print "Error: " + str(e)
		return 1
	return 0

if __name__ == "__main__":
	sys.exit(main())