#!/usr/bin/python
#
# A command-line LiveJournal friendslist manipulation program.
#
# Available operations:
#	add	- add new friends
#	grm	- remove friends from groups
#	groups	- list the names of your friend groups
#	ln	- add friends to groups
#	ls	- list friends, with what friend groups they are in.
#	lsrev	- list people who friend you that you do not friend.
#	mv	- move friends out of friends groups and into others.
#	neaten	- try to get everyone into just one friend group.
#	remove	- remove friends.
#
# Your LiveJournal userid and password is expected to be in $HOME/.ljfrc
# in a line of the form 'login <user> <password>'. Blank lines and comment
# lines (starting with a '#' as the first non-whitespace) are allowed.
#
# grm and mv can take an augmented friend-naming syntax of the form
# 'group/friend'; this means 'apply the operation to just this group
# that the friend is in'. Otherwise, the operation is applied to all
# groups that the friend is in. So 'mv silly/fred normal' means to
# remove fred from silly and add fred to normal; 'mv fred normal'
# means add fred to normal, removing him from all other
# groups. Similarly 'grm silly/fred' means 'remove fred from silly',
# whereas 'grm fred' means 'remove fred from all friends groups'.
#
# (Technically ln also takes this syntax, but it does nothing
# useful in ln's case; 'ln silly/fred normal' and 'ln fred normal'
# do the exact same thing.)
#
# neaten does the following two things:
# 1: if someone is in no friends groups, they are added to Default View
# 2: if someone is in more than one group, and Default View is one of them,
#    they are removed from Default View.
# Neaten does nothing if you have no Default View group.
# Neaten always reports what it is doing; you cannot shut it up. -n makes
# it 'dryrun', not actually commit the changes.
# Other operations are generally silent.
#
# ls options:
# -[012]	print (only) people in no friends groups, one friend group,
#		or more than one friend group respectively. normally
#		everyone is printed, which is equivalent to -012
# -q		never print what groups people are in
# -v		always print what groups people are in
#		by default, ls prints what groups people are in unless
#		you are listing just one group.
# -x		do not list people who are only in Default View.
#
# lsrev options:
# -a		list all people who friend you
# -m		list only mutual friends.
# (so it is sort of a general friendof lister now)

import sys, os, os.path
import xmlrpclib, socket, md5, time
import sets, getopt

class LJChatter:
	"""Perform authenticated XMLRC with LiveJournal."""
	def __init__(self, user, pw):
		self._u = user
		self._hp = md5.new(pw).hexdigest()
		# Obtain the starting XMLRPC handle.
		self.srv = xmlrpclib.ServerProxy("http://www.livejournal.com/interface/xmlrpc").LJ.XMLRPC

	def withauth(self, ljcall, data):
		"""Perform a LiveJournal call with challenge/response
		authentication.

		The ljcall argument must be a ready to go XMLRPC
		handle to the desired XMLRPC operation."""		
		res = self.srv.getchallenge()
		ch = res["challenge"]
		data["auth_method"] = "challenge"
		data["auth_challenge"] = ch
		data["auth_response"] = md5.new("%s%s" % (ch, self._hp)).hexdigest()
		data["username"] = self._u
		data["ver"] = 1

		return ljcall(data)

	def getfriends(self, lsrev = False):
		"""Obtain the LiveJournal friendslist, optionally including
		friendsof data."""
		data = {"includegroups": 1}
		if lsrev:
			data["includefriendof"] = 1
		return self.withauth(self.srv.getfriends, data)

	def mvfriendgroups(self, trans):
		"""Change the friends groups that friends are in, possibly
		adding friends in the process.

		(This uses LJ.XMPLRC.editfriends, not .editfriendsgroups.)"""
		return self.withauth(self.srv.editfriends, trans.to_data())
	def delfriends(self, trans):
		"""Unfriend people."""
		return self.withauth(self.srv.editfriends, trans.to_del_data())

# This code assumes that people will not make contradictory changes.
NGRPS = 1
class LJTrans:
	"""Organize and track a series of changes to LiveJournal friends."""
	def __init__(self, groups):
		self.groups = groups
		self.modusers = {}
		self.delusers = {}

	def _setup(self, user):
		un = user['username']
		if un not in self.modusers:
			self.modusers[un] = user.get("groupmask", NGRPS)

	def clear_groups(self, user):
		"""Remove a friend from all friends groups."""
		self.modusers[user['username']] = NGRPS
	def add_group(self, user, groupname):
		"""Add a friend to a friend group."""
		self._setup(user)
		gm = self.groups.get_mask(groupname)
		self.modusers[user['username']] |= gm
	def del_group(self, user, groupname):
		"""Remove a friend from a friend group."""
		self._setup(user)
		gm = self.groups.get_mask(groupname)
		self.modusers[user['username']] &= ~gm

	def add_user(self, username, group = 'Default View'):
		"""Add a new friend."""
		gm = self.groups.get_mask(group)
		if not gm:
			gm = NGRPS
		self.modusers[username] = gm

	def del_user(self, username):
		"""Remove a friend."""
		self.delusers[username] = True

	def to_data(self):
		"""Compile friend additions and friend group changes into
		the LJ XMLRPC format."""
		assert not self.delusers, "self.delusers set in LJTrans.to_data: "+str(self.delusers)
		changed = []
		for k, v in self.modusers.iteritems():
			changed.append({'username': k, 'groupmask': v})
		return {'add': changed}

	def to_del_data(self):
		"""Compile friend deletions into the LJ XMLRPC format."""
		assert not self.modusers, "self.modusers set in LJTrans.to_del_data: "+str(self.modusers)
		return {'delete': self.delusers.keys()}

	def __len__(self):
		return len(self.modusers) + len(self.delusers)


class LJFError(Exception):
	pass
class LJFUsage(Exception):
	pass
class LJFArgs(Exception):
	pass

class Groups:
	"""Store and index and process the list of LiveJournal friendsgroups."""
	def __init__(self, fgraw):
		self.nameidx = {}
		self.maskidx = []
		self.gCache = {NGRPS: ()}
		for fg in fgraw:
			name = fg['name']
			idn = fg['id']
			self.nameidx[name] = fg
			self.maskidx.append((1 << idn, name))
			# We might as well preload the cache.
			self.gCache[(1 << idn) | 1] = (name,)

	def get_groups(self, friend):
		"""Get a list of the group names that a friend entry is in."""
		gmask = friend.get("groupmask", NGRPS)
		if gmask not in self.gCache:
			gl = [x[1] for x in self.maskidx if x[0] & gmask]
			gl.sort()
			gl = tuple(gl)
			self.gCache[gmask] = gl
		return self.gCache[gmask]

	def get_mask(self, groupname):
		"""Get the groupmask for a particular friend group."""
		if groupname not in self.nameidx:
			return 0
		return (1 << self.nameidx[groupname]['id'])

	def __getattr__(self, name):
		return getattr(self.nameidx, name)
	def keys(self):
		return self.nameidx.keys()

class Friends:
	"""Store and index the list of LiveJournal friends."""
	def __init__(self, flist):
		self.fidx = {}
		for f in flist:
			self.fidx[f['username']] = f
	def __getattr__(self, name):
		return getattr(self.fidx, name)
	def keys(self):
		return self.fidx.keys()

def get_friendslist(cxn, lsrev = False):
	"""Retrieve the LiveJournal friendslist and parse it into
	structures."""
	# Every now and then LiveJournal freaks out and drops the
	# friendslist (although not the groups list). Go figure.
	# So we retry it a couple of times.
	cnt = 3
	while cnt:
		r = cxn.getfriends(lsrev)
		if 'friends' in r and len(r['friends']) > 0:
			break
		cnt = cnt - 1
		sys.stderr.write("%s: LiveJournal XMLRPC problem, retrying: no friendslist returned.\n" % sys.argv[0])
		time.sleep(2)
		continue

	# Hopefully we have gotten something back.
	friends = Friends(r['friends'])
	groups = Groups(r['friendgroups'])
	return (friends, groups, r)

def setup_trans(cxn, musthave = None):
	"""Set up for performing a group/friendslist change transaction.

	Get the list of friends and groups from LiveJournal, check for
	an optional group that must be in the grouplist, and return all
	this plus a transaction object."""
	friends, groups, r = get_friendslist(cxn)
	if musthave and musthave not in groups:
		raise LJFArgs, "No such friends group: "+musthave
	return friends, groups, r, LJTrans(groups)

def pretty_group_name(name):
	"""Pretty up the name of a friendsgroup by putting quotes around
	names with spaces in them."""
	if ' ' in name:
		return "'%s'" % name
	else:
		return name
def pretty_grouplist(gl):
	"""Format a grouplist tuple for printing"""
	return ', '.join([pretty_group_name(x) for x in gl])

def cmd_ls(cxn, opts, args):
	"Implement the ls command"
	lst_all = True
	lst_none = False
	lst_one = False
	lst_mult = False
	no_defview = False
	quiet = False
	verbose = False
	for o, a in opts:
		if o == '-0':
			lst_all = False
			lst_none = True
		elif o == '-1':
			lst_all = False
			lst_one = True
		elif o == '-2':
			lst_all = False
			lst_mult = True
		elif o == '-q':
			quiet = True
		elif o == '-v':
			verbose = True
		elif o == '-x':
			no_defview = True
		else:
			raise AssertionError, "unhandled ls option: "+o
	if quiet and verbose:
		raise LJFArgs, "cannot use both -q and -v"

	friends, groups, r = get_friendslist(cxn)
	fk = friends.keys()
	fk.sort()
	if not fk:
		print "No friends returned in the result!"
		# this is debugging output, and I should do better.
		print r
		return

	select_set = sets.ImmutableSet(args)
	for fn in fk:
		f = friends[fn]
		gl = groups.get_groups(f)

		# Filter out conditional groups.
		if not lst_all:
			if lst_none and len(gl) == 0:
				pass
			elif lst_one and len(gl) == 1:
				pass
			elif lst_mult and len(gl) > 1:
				pass
			else:
				continue

		# Exclude people in only the Default View, if asked.
		if no_defview and len(gl) == 1 and gl[0] == 'Default View':
			continue

		# Are we restricted to printing members of just the
		# groups that we got as arguments?
		if select_set:
			gset = sets.ImmutableSet(gl)
			if not gset.intersection(select_set):
				continue

		# Do we print the groups?
		if quiet or (len(args) == 1 and not verbose):
			grps = ''
		else:
			grps = pretty_grouplist(gl)
		print "%-20s %s" % (fn, grps)

def cmd_lsrev(cxn, opts, args):
	"Implement the lsrev command"
	__pychecker__ = "no-argsused"
	show_all = False
	show_mutual = False
	for o, a in opts:
		if o == '-a':
			show_all = True
		elif o == '-m':
			show_mutual = True
		else:
			raise AssertionError, "unhandled lsrev option: "+o
	if show_all and show_mutual:
		raise LJFArgs, "-a and -m cannot be combined"
	friends, groups, r = get_friendslist(cxn, True)
	if 'friendofs' not in r:
		print "%s lsrev: could not get the friendof list" % sys.argv[0]
		return
	for f in r['friendofs']:
		mutual = f['username'] in friends
		if show_all:
			pass
		# This is the rare case where a logical XOR would be useful.
		elif (show_mutual and not mutual) \
		     or (not show_mutual and mutual):
			continue
		print "%-20s" % f['username']

def cmd_lsgrps(cxn, opts, args):
	"Implement the groups command"
	__pychecker__ = "no-argsused"
	_, groups, _ = get_friendslist(cxn)
	k = groups.keys()
	k.sort()
	for g in k:
		print g

def cmd_neaten(cxn, opts, args):
	"Implement the neaten command"
	__pychecker__ = "no-argsused"	
	dryrun = False
	for o, a in opts:
		if o == "-n":
			dryrun = True
		else:
			raise AssertionError, "unhandled neaten option: "+o

	# We can't do this if there's no 'Default View' in groups.
	friends, groups, _, trans = setup_trans(cxn, 'Default View')

	fk = friends.keys()
	fk.sort()
	for fn in fk:
		f = friends[fn]
		gl = groups.get_groups(f)
		if len(gl) == 0:
			print "%-20s add Default View: in no groups" % fn
			trans.add_group(f, 'Default View')
		elif len(gl) > 1 and 'Default View' in gl:
			print "%-20s del Default View: in %s" % (fn, pretty_grouplist(gl))
			trans.del_group(f, 'Default View')
	if trans and not dryrun:
		print "%s neaten: comitting changes" % sys.argv[0]
		cxn.mvfriendgroups(trans)

def argsplit(arg):
	"""Split a possible Group/friend argument into group and friend.
	None is returned for anything missing."""
	if '/' not in arg or arg[-1] == '/':
		return (None, arg)
	r = arg.rfind("/")
	return (arg[:r], arg[r+1:])

def walkfriends(friends, args):
	"""Walk a list of friendsblobs, splitting into group and username,
	check to see if they are in friends, and yield the split up stuff."""
	for fstr in args:
		(grp, fn) = argsplit(fstr)
		# I chose not to implicitly add friends in this operation,
		# because this way I guard against my own stupid typos.
		if fn not in friends:
			raise LJFArgs, "No such friend: "+fn
		f = friends[fn]
		yield grp, fn, f

# Much of the code is common over ln and mv; the only difference
# is whether people are moved out of their old group(s) when we
# put them in a new one.
def mvln_common(cxn, args, del_old):
	"""The common core of mv and ln"""
	tgrp = args[-1]
	if tgrp[-1] == '/':
		tgrp = tgrp[:-1]
	frnds = args[:-1]

	friends, groups, _, trans = setup_trans(cxn, tgrp)
	for grp, fn, f in walkfriends(friends, frnds):
		if del_old:
			if not grp:
				trans.clear_groups(f)
			else:
				trans.del_group(f, grp)
		trans.add_group(f, tgrp)
	cxn.mvfriendgroups(trans)
	
def cmd_mv(cxn, opts, args):
	"Implement the mv command"
	__pychecker__ = "no-argsused"
	mvln_common(cxn, args, True)

def cmd_ln(cxn, opts, args):
	"Implement the ln command"
	__pychecker__ = "no-argsused"
	mvln_common(cxn, args, False)

def cmd_grm(cxn, opts, args):
	"Implement the grm command"
	__pychecker__ = "no-argsused"
	friends, groups, _, trans = setup_trans(cxn)
	for grp, fn, f in walkfriends(friends, args):
		if not grp:
			trans.clear_groups(f)
		else:
			trans.del_group(f, grp)
	cxn.mvfriendgroups(trans)

def cmd_friend(cxn, opts, args):
	"Implement the add command"
	tgrp = None
	for o, a in opts:
		if o == '-g':
			tgrp = a
		else:
			raise AssertionError, "unhandled friend option: "+o
	friends, groups, _, trans = setup_trans(cxn, tgrp)
	if not tgrp:
		tgrp = 'Default View'
	for fstr in args:
		if fstr in friends:
			print "%-20s already a friend" % fstr
		else:
			trans.add_user(fstr, tgrp)
	cxn.mvfriendgroups(trans)

def cmd_unfriend(cxn, opts, args):
	"Implement the remove command"
	__pychecker__ = "no-argsused"
	friends, groups, _, trans = setup_trans(cxn)
	for fstr in args:
		if fstr not in friends:
			raise LJFArgs, "no friend '%s'" % fstr
		trans.del_user(fstr)
	cxn.delfriends(trans)

def die(msg):
	"Exit with an error message."
	sys.stderr.write("%s: %s\n" % (sys.argv[0], msg))
	sys.exit(1)

def get_user_pass():
	"""Read the login and password from $HOME/.ljfrc."""
	if 'HOME' not in os.environ:
		raise LJFError, "no $HOME environment variable"
	rcfile = os.path.join(os.environ["HOME"], ".ljfrc")
	try:
		fp = open(rcfile, "r")
		while 1:
			l = fp.readline()
			if not l:
				break
			l = l.strip()
			if not l or l[0] == "#":
				continue
			t = l.strip().split(None, 2)
			if len(t) != 3 or t[0] != "login":
				continue
			else:
				return t[1], t[2]
		raise LJFError, "could not find login information in "+rcfile
	except EnvironmentError, e:
		raise LJFError, "problem reading %s: %s" % (rcfile, str(e))

def startup():
	"""Start up a LiveJournal connection object. Fishes user and
	password information from $HOME/.ljfrc."""
	try:
		r = get_user_pass()
	except LJFError, e:
		die("could not get user and password: "+str(e))

	return LJChatter(r[0], r[1])

def usage():
	"Report ljf's usage message."
	sys.stderr.write("usage: %s <cmd> [cmd args]\n" % sys.argv[0])
	sys.stderr.write("available commands and arguments:\n")
	kys = cmds.keys()
	kys.sort()
	for k in kys:
		sys.stderr.write("    %s %s\n" % (k, cmds[k][1]))
	sys.exit(1)

cmds = {'ls': (cmd_ls, '[-012] [-q|v] [-x] [group [group ...]]',
	       "012qvx", 0, -1),
	'neaten': (cmd_neaten, "[-n]", "n", 0, 0),
	"mv": (cmd_mv, "friend [friend ...] group", "", 2, -1),
	"ln": (cmd_ln, "friend [friend ...] group", "", 2, -1),
	"grm": (cmd_grm, "friend [friend ...]", "", 1, -1),
	"add": (cmd_friend, "[-g group] friend [friend ...]", "g:", 1, -1),
	"remove": (cmd_unfriend, "friend [friend ...]", "", 1, -1),
	"lsrev": (cmd_lsrev, "[-a|m]", "am", 0, 0),
	"groups": (cmd_lsgrps, "", "", 0, 0),
	}

def main(args):
	"Perform main ljf processing."
	if len(args) < 1:
		usage()
	cxn = startup()

	# Pick function
	cmd = args[0]
	if cmd in cmds:
		(cmd_func, ustr, gostr, mina, maxa) = cmds[cmd]
	else:
		usage()
	args = args[1:]

	# Run the function.
	try:
		opts, args = getopt.getopt(args, gostr)
		if len(args) < mina:
			raise LJFUsage, "too few arguments"
		if maxa == 0 and args:
			raise LJFUsage, "takes no arguments"
		cmd_func(cxn, opts, args)
	except xmlrpclib.Fault, e:
		die("problem talking with LiveJournal: "+e.faultString)
	except xmlrpclib.ProtocolError, e:
		die("XMLPRC problem: "+str(e))
	except socket.error, e:
		die("network problem while talking with LiveJournal: "+str(e))
	except (getopt.GetoptError, LJFUsage), e:
		sys.stderr.write("%s %s: %s\n" % (sys.argv[0], cmd, e))
		sys.stderr.write("usage %s %s %s\n" % (sys.argv[0], cmd, ustr))
		sys.exit(1)
	except LJFArgs, e:
		sys.stderr.write("%s %s: %s\n" % (sys.argv[0], cmd, e))
		sys.exit(1)

if __name__ == "__main__":
	main(sys.argv[1:])