import-vobjects-to-sogo

#!/usr/bin/python3

# Script for add contacts, calendars and todos from files SOGo database
# (c) 2016 Andrey Cherepanov <cas@altlinux.org>

VERSION = '1.1.3'
verbose = False
folders = []

import argparse
import sys
import os
import psycopg2
import psycopg2.errorcodes
import vobject
import uuid
import hashlib
import time
from pprint import pprint
from psycopg2.extras import NamedTupleCursor
from datetime import datetime, timedelta

# Supported types
supported_types = {
	 'vcard':  'contacts',
	 'calendar': 'calendar',
	}
# Default folder names
default_folders = [ 'Контакты', 'Contacts', 'Календарь', 'Calendar' ]

# Get timestamp
timestamp = int(time.time())

# Functions
def generate_uuid( len=24 ):
	"""Generate filder UUID"""
	m = hashlib.sha224()
	m.update( repr( uuid.uuid4() ).encode( 'utf-8' ) )
	z = m.digest().hex()
	u = ( z[0:4] + '-' + z[4:12] + '-' + z[13] + '-' + z[14:21] ).upper()
	if len == 24:
		return u
	else:
		return z[:len]

# Classes
class DB:
	"""Class for database connection"""
	conn = None
	cursor = None
	ncursor = None
	connection_string = ''
	location = ''
	user = ''
	last_query = ''

	def __init__( self ):
		"""Init connection"""

	def connect( self, connection_string, user_name, location='' ):
		"""Connect to database"""
		self.connection_string = connection_string
		self.user = user_name
		self.location = location
		if self.location == '':
			self.location = self.connection_string
		if verbose:
			print( "Connect to database %s to serve user %s..." % ( self.connection_string, self.user ) )
		self.conn = psycopg2.connect( self.connection_string )
		self.cursor = self.conn.cursor()
		self.ncursor = self.conn.cursor( cursor_factory=NamedTupleCursor )

	def query( self, query, args=None, as_dict=False, debug=False ):
		"""Perform SQL query: as string or with optional arguments"""
		if as_dict == True:
			# Return as dictionary
			records = None
			try:
				if args != None:
					self.last_query = self.ncursor.mogrify( query, args )
				else:
					self.last_query = query
				if debug:
					print( self.last_query )
				res = self.ncursor.execute( self.last_query )
				records = self.ncursor.fetchone()
			except psycopg2.ProgrammingError:
				pass
#			except psycopg2.Error as e:
#				print( "DATABASE ERROR [%s] %s" % ( e.pgcode, e.diag.message_primary ) )
#				print(psycopg2.errorcodes.lookup(e.pgcode[:2]))
#				print(psycopg2.errorcodes.lookup(e.pgcode))
		else:
			# Ordinary cursor
			records = []
			try:
				if args != None:
					self.last_query = self.cursor.mogrify( query, args )
				else:
					self.last_query = query
				if debug:
					print( self.last_query )
				res = self.cursor.execute( self.last_query )
				if debug:
					print( 'Result: ', self.cursor.rowcount )
				if self.cursor.rowcount > 0:
					records = self.cursor.fetchall()
				if self.cursor.rowcount == -1 and self.last_query[:12] != 'CREATE TABLE' and self.last_query != 'COMMIT':
					print( "Error in query %s" % ( self.last_query ) )
			except psycopg2.Error as e:
				if e.pgcode != None:
					print( "DATABASE ERROR [%s] %s" % ( e.pgcode, e.diag.message_primary ) )
					print(psycopg2.errorcodes.lookup(e.pgcode[:2]))
					print(psycopg2.errorcodes.lookup(e.pgcode))
			except:
				#print( "Error in query %s" % ( self.cursor.query ) )
				raise
		return records

	def disconnect( self ):
		"""Disconnect from database"""
		if verbose:
			print( "Disconnect from database %s..." % ( self.connection_string ) )
		self.conn.close()

db = DB()

class DBFolder:
	"""Class for store Folder in database"""
	type = ''
	user = ''
	name = ''
	real_name = ''
	table_name = ''
	c_folder_id = ''
	c_path = ''
	c_path1 = 'Users'
	c_path2 = ''
	c_path3 = ''
	c_path4 = 'personal'
	c_foldername = ''
	c_location = ''
	c_quick_location = ''
	c_acl_location = ''
	c_folder_type = ''
	uuid = ''
	short_uuid = ''
	db_main_table = ''

	def __init__( self, type, user, name ):
		"""Init values for user and folder name: use existing of fill defaults"""
		self.type = type
		self.user = user
		self.name = name

		if type == 'vcard':
			self.c_path3 = 'Contacts'
			self.c_folder_type = 'Contact'
		if type == 'calendar':
			self.c_path3 = 'Calendar'
			self.c_folder_type = 'Appointment'

		# Process folder name and path
		self.c_path2 = self.user
		base_name  = name.split( '/' )[-1:][0]
		if not base_name in default_folders:
			self.uuid = generate_uuid()
			self.c_path4 = self.uuid
		else:
			self.c_path4 = 'personal'
		self.c_foldername = base_name
		self.c_path = '/' + self.c_path1 + '/' + self.c_path2 + '/'+ self.c_path3 + '/'+ self.c_path4

		# Short UUID and locations
		self.short_uuid = generate_uuid( 11 )
		# Example: postgresql://sogo:1234@localhost:5432/sogo/sogodod001541a8862
		self.db_main_table = 'sogo' + db.user + self.short_uuid
		self.c_location = db.location + '/' + self.db_main_table
		self.c_quick_location = self.c_location + '_quick'
		self.c_acl_location = self.c_location + '_acl'

	def save( self ):
		"""Save FolderRecord to database table sogo_folder_info, create tables for folder data"""
		if verbose:
			print( "Check if %s exists" % ( self.c_path ) )
		# Check if record exists
		res = db.query( "select * from sogo_folder_info where c_folder_type=%s and c_path2=%s and c_foldername=%s;",
				[ self.c_folder_type, self.c_path2, self.c_foldername ],
				as_dict=True )

		if res != None:
			# Record is exist
			self.c_path  = res.c_path
			self.c_path1 = res.c_path1
			self.c_path2 = res.c_path2
			self.c_path3 = res.c_path3
			self.c_path4 = res.c_path4
			self.c_foldername = res.c_foldername
			self.c_location = res.c_location
			self.c_quick_location = res.c_quick_location
			self.c_acl_location = res.c_acl_location
			self.c_folder_type = res.c_folder_type
			self.db_main_table = self.c_location.split( '/' )[-1:][0]
		else:
			# New record
			if verbose:
				print( "Add folder: %s" % ( self.c_path ) )
			db.query( """insert into sogo_folder_info(c_path,c_path1,c_path2,c_path3,c_path4,c_foldername,c_location,c_quick_location,c_acl_location,c_folder_type)
			           values(%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
			           [ self.c_path, self.c_path1, self.c_path2, self.c_path3, self.c_path4, self.c_foldername,
			             self.c_location, self.c_quick_location, self.c_acl_location, self.c_folder_type ] )

		# Create data table for folder
		if verbose:
			print( "CREATE TABLE %s" % ( self.db_main_table ) )
			print( "CREATE TABLE %s_quick" % ( self.db_main_table ) )
			print( "CREATE TABLE %s_acl" % ( self.db_main_table ) )

		db.query( """CREATE TABLE IF NOT EXISTS %s (
				c_name varchar(255) NOT NULL,
				c_content text NOT NULL,
				c_creationdate int4 NOT NULL,
				c_lastmodified int4 NOT NULL,
				c_version int4 NOT NULL,
				c_deleted int4 NULL,
				CONSTRAINT %s PRIMARY KEY (c_name) );""" %
				( self.db_main_table, self.db_main_table+'_pkey' ) )
		if self.type == 'vcard':
			db.query( """CREATE TABLE IF NOT EXISTS %s (
					c_name varchar(255) NOT NULL,
					c_givenname varchar(255) NULL,
					c_cn varchar(255) NULL,
					c_sn varchar(255) NULL,
					c_screenname varchar(255) NULL,
					c_l varchar(255) NULL,
					c_mail varchar(255) NULL,
					c_o varchar(255) NULL,
					c_ou varchar(255) NULL,
					c_telephonenumber varchar(255) NULL,
					c_categories varchar(255) NULL,
					c_component varchar(10) NOT NULL,
					CONSTRAINT %s PRIMARY KEY (c_name) );""" %
					( self.db_main_table+'_quick', self.db_main_table+'_quick_pkey' ) )
		elif self.type == 'calendar':
			db.query( """CREATE TABLE IF NOT EXISTS %s (
					c_name varchar(255) NOT NULL,
					c_uid varchar(255) NOT NULL,
					c_startdate int4 NULL,
					c_enddate int4 NULL,
					c_cycleenddate int4 NULL,
					c_title varchar(1000) NOT NULL,
					c_participants text NULL,
					c_isallday int4 NULL,
					c_iscycle int4 NULL,
					c_cycleinfo text NULL,
					c_classification int4 NOT NULL,
					c_isopaque int4 NOT NULL,
					c_status int4 NOT NULL,
					c_priority int4 NULL,
					c_location varchar(255) NULL,
					c_orgmail varchar(255) NULL,
					c_partmails text NULL,
					c_partstates text NULL,
					c_category varchar(255) NULL,
					c_sequence int4 NULL,
					c_component varchar(10) NOT NULL,
					c_nextalarm int4 NULL,
					c_description text NULL,
					CONSTRAINT %s PRIMARY KEY (c_name) );""" %
					( self.db_main_table+'_quick', self.db_main_table+'_quick_pkey' ) )
		db.query( """CREATE TABLE IF NOT EXISTS %s (
				c_uid varchar(255) NOT NULL,
				c_object varchar(255) NOT NULL,
				c_role varchar(80) NOT NULL );""" %
				( self.db_main_table+'_acl' ) )
		db.query( 'COMMIT' )
		return self.db_main_table

class vCardRecord:
	"""Class of single vCard"""
	obj = None
	file = ''
	table = ''

	def __init__( self, obj ):
		"""Create new vCard object to put into database"""
		self.obj = obj
		self.file = generate_uuid() + '.vcf'

	def saveCard( self ):
		"""Save vCard content to table sogo<user><uuid>"""
		try:
			db.query( """insert into """ + self.table + """(c_name,c_content,c_creationdate,c_lastmodified,c_version)
				  values(%s,%s,%s,%s,%s)""",
				  [ self.file, self.obj.serialize(), timestamp, timestamp, 0 ] )
		except psycopg2.ProgrammingError:
			pass
		except:
			raise

	def saveValues( self ):
		"""Save vCard values to table sogo<user><uuid>_quick"""

		# Prepare list of columns of existing values
		columns = [ 'name', 'component' ]
		values  = [ self.file, 'vcard' ]
		types = []

		for key, value in {
			'givenname': self.obj.n.value.given,
			'cn': self.obj.fn.value if self.obj.fn.value != '(null)' else None,
			'sn': self.obj.n.value.family,
			'screenname': self.obj.getChildValue( 'X-AIM' ),
			'l': self.obj.adr.value.city if self.obj.getChildValue( 'adr' ) != None else None,
			'mail': self.obj.getChildValue( 'email' ),
			'o': self.obj.org.value[0] if self.obj.getChildValue('org')!=None else None,
			'ou': ', '.join(self.obj.org.value[1:]) if self.obj.getChildValue('org')!=None else None,
			'telephonenumber': self.obj.getChildValue( 'tel' ),
			'categories': self.obj.getChildValue( 'categories' ),
		}.items():
			if value != None:
				columns.append( key )
				values.append( value )

		# Modify lists
		for i, col in enumerate( columns ):
			columns[i] = 'c_' + columns[i]
			types.append( '%s' )

		try:
			db.query( 'insert into %s_quick(%s) values(%s)' %
				  ( self.table, ','.join( columns ), ','.join( types ) ),
				  values )
		except psycopg2.Error as e:
			if e.pgcode == None:
				pass
			else:
				raise

	def save( self, table='' ):
		"""Save vCard data in two tables"""
		if table != '':
			self.table = table
			self.saveCard()
			self.saveValues()

class calendarRecord:
	"""Class of single calendar entry"""
	type = ''
	ev = None
	obj = None
	file = ''
	table = ''

	def __init__( self, obj ):
		"""Create new vCard object to put into database"""
		self.obj = obj
		self.file = generate_uuid() + '.ics'
		if 'vevent' in self.obj.contents:
			self.type = 'vevent'
			self.ev = self.obj.vevent
		else:
			self.type = 'vtodo'
			self.ev = self.obj.vtodo
		if self.ev.getChildValue( 'uid' ) != None:
			self.file = self.ev.uid.value

	def saveCard( self ):
		"""Save vCard content to table sogo<user><uuid>"""
		try:
			db.query( "delete from " + self.table + " where c_name='" + self.file + "'" )
			db.query( """insert into """ + self.table + """(c_name,c_content,c_creationdate,c_lastmodified,c_version)
				  values(%s,%s,%s,%s,%s)""",
				  [ self.file, self.obj.serialize(), timestamp, timestamp, 0 ] )
		except psycopg2.ProgrammingError:
			pass
		except:
			raise

	def saveValues( self ):
		"""Save vCard values to table sogo<user><uuid>_quick"""

		if self.ev == None:
			return

		# Prepare list of columns of existing values
		columns = [ 'name' ]
		values  = [ self.file ]
		types = []
		cycle_end_date = None
		status = 0

		ev = self.ev

		# Some fields computation
		rrule = ev.getChildValue( 'rrule' )
		if rrule != None:
			# Example: UNTIL=20160815T210000Z
			for param in rrule.value.split( ';' ):
				if param[:6] == 'UNTIL=':
					cycle_end_date = int(datetime.strptime(param[6:],'%Y%m%dT%H%M%SZ').strftime('%s'))-time.timezone
		if ev.getChildValue( 'status' ) == 'CONFIRMED':
			status = 1

		# Fill fields
		data = {
			'uid': ev.uid.value if ev.getChildValue( 'uid' ) != None else self.file[:-4],
			'startdate': 0 if ev.getChildValue( 'dtstart' ) == None else int(ev.dtstart.value.strftime('%s'))-time.timezone,
			'enddate': 0 if ev.getChildValue( 'dtend' ) == None else int(ev.dtend.value.strftime('%s'))-time.timezone,
			'cycleenddate': cycle_end_date,
			'title': ev.getChildValue( 'summary', '' ),
			'participants': None if ev.getChildValue('attendee') == None else ','.join( ev.attendee.params.get('CN') ),
			'isallday': 1 if ev.getChildValue( 'dtstart' ) != None and ev.getChildValue( 'dtend' ) != None and ev.dtend.value-ev.dtstart.value == timedelta(1) else 0,
			'iscycle': 0 if rrule == None else 1,
			'cycleinfo': None if rrule == None else '{rules = ("%s"); }' % ( rrule.value ),
			'classification': 0, # TODO ?
			'isopaque': 0 if ev.getChildValue( 'transp' ) != None and ev.getChildValue( 'transp' ) != 'OPAQUE' else 1,
			'status': status,
			'priority': 0, # TODO ?
			'location': ev.getChildValue( 'location' ),
			'orgmail': ev.getChildValue( 'organizer', '' ),
			'partmails': ev.getChildValue( 'attendee' ),
			'partstates': '', # TODO: ev.attendee.params.get('PARTSTAT') = ['NEEDS-ACTION']
			'category': None if ev.getChildValue( 'categories' ) == None else ','.join(ev.categories.value),
			'sequence': 0 if ev.getChildValue( 'sequence' ) == None else int(ev.sequence),
			'component': self.type,
			'nextalarm': 0, # TODO need algorithm for calc next alarm
			'description': ev.getChildValue( 'description', '' )
		}

		for key, value in data.items():
			if value != None:
				columns.append( key )
				values.append( value )

		# Modify lists
		for i, col in enumerate( columns ):
			columns[i] = 'c_' + columns[i]
			types.append( '%s' )

		try:
			db.query( "delete from " + self.table + "_quick where c_name='" + self.file + "'" )
			db.query( 'insert into %s_quick(%s) values(%s)' %
				  ( self.table, ','.join( columns ), ','.join( types ) ),
				  values )
		except psycopg2.Error as e:
			if e.pgcode == None:
				pass
			else:
				raise

	def save( self, table='' ):
		"""Save vCard data in two tables"""
		if table != '':
			self.table = table
			self.saveCard()
			self.saveValues()

class Folder:
	"""Class of vobject folder"""
	type = ''
	items = []
	file_name = ''
	folder_path = []
	record = None
	db_main_table = ''

	def __init__( self, type, base_path, directory ):
		"""Create new folder instance"""
		self.type = type
		if type in supported_types:
			self.file_name = os.path.join( base_path, directory, supported_types[ type ] )
			self.folder_path = directory.split( os.sep )
		self.items = []
		db_main_table = ''

	def read( self ):
		"""Open file and read all vobject items in self.items"""
		if verbose:
			print( "Folder.read(%s)..." % ( self.file_name ) )
		try:
			stream = open( self.file_name, 'r', encoding='utf-8' )
			for obj in vobject.readComponents( stream ):
				if self.type == 'vcard' and obj.name == 'VCARD':
					# Put in list vCardRecord object
					self.items.append( vCardRecord( obj ) )
				if self.type == 'calendar' and obj.name == 'VCALENDAR':
					# Put in list calendarRecord object
					self.items.append( calendarRecord( obj ) )
			if verbose:
				print( "Read records: %d" % ( len( self.items ) ) )
			stream.close()
		except:
			print( "Error reading records from %s" % ( self.file_name ) )
			raise

	def save( self, cursor ):
		"""Store values to specied database"""
		if verbose:
			print( "Folder.save(%s)..." % ( db.connection_string ) )
		# Create folders structure
		if verbose:
			print( "Create folder in database: %s" % ( os.sep.join( self.folder_path ) ) )
		self.record = DBFolder( self.type, db.user, os.sep.join( self.folder_path ) )
		self.db_main_table = self.record.save()

		# Dump vCard contents
		if verbose:
			print( "Folder.save.items" )
		for rec in self.items:
			# Store item content and some values to database
			rec.save( self.db_main_table )

# Parse command line arguments
parser = argparse.ArgumentParser( description='Script for import contacts, calendars and todos from files to SOGo database',
                                  epilog="Example: ./import-vobjects-to-sogo --db 'postgresql://sogo@/sogo/sogo_folder_info' dod" )
parser.add_argument( 'user_name', help='SOGo user name' )
parser.add_argument( '--db', dest='db_connect', required=True,
                     help='SOGo database connection string (see OCSFolderInfoURL in /etc/sogo/sogo.conf)' )
parser.add_argument( '--location', dest='location',
                     help='SOGo location string (see OCSFolderInfoURL in /etc/sogo/sogo.conf without folder)' )
parser.add_argument( '--dir', dest='directory', default='.', help='Path to directory with imported folders' )
parser.add_argument( '-v', dest='verbose', action='count', help='Be more verbose' )
parser.add_argument( '--version', action='version', version=VERSION )
args = parser.parse_args()

if args.verbose != None:
	verbose = True
	print( 'Run with following arguments:' )
	pprint( args )

# Connect to database
try:
	# Get a connection
	db.connect( args.db_connect, args.user_name, args.location )
except:
	print( "Cannot connect to database %s:" % ( args.db_connect ) )
	raise

# Walk through folders containing file contacts
for cdir, dirs, files in os.walk( args.directory ):
	for type, file_name in supported_types.items():
		if file_name in files:
			print( "Folder %s/%s (%s)" % ( cdir, file_name, type ) )
			# Remove basedir from path and possible os.sep
			if args.directory[-1:] == os.sep:
				args.directory = args.directory[:-1]
			folder = Folder( type, args.directory, cdir[len(args.directory)+1:] )
			folders.append( folder )
			# Open file and store as list of vobject objects (vCard)
			folder.read()
			# Store in database
			folder.save( args.db_connect )