#!/usr/bin/python

import os,sys
import json
import glob
import subprocess

import rados
import rbd

import pwd, grp

def print_help():
	print "Usage: %s <conffile> <command> <slot> <path> <driveindex>"%(sys.argv[0])

class RBDChanger(object):
	def __init__(self, conffile):
		self.conffile = conffile
		self.load_config()
	
	def load_config(self):
		d = json.loads(open(self.conffile).read() )
		self.config = d
		#print d

	def __run_command(self, cmd):
		p1 = subprocess.Popen(["sudo"] + cmd, stdout=subprocess.PIPE)
		p1.wait()
		
		output = p1.communicate()
		#print output
		return (p1.returncode, output[0])

	def __getmapped(self):
		cmd = [
			"/usr/bin/rbd",
			"showmapped",
			"--pool", self.config['pool'],
			'--format','json'
		]
		#print cmd
		status,output = self.__run_command(cmd)
		j = json.loads(output)
		return j

	def __open_rados(self):
		rconn = rados.Rados(conffile='/etc/ceph/ceph.conf')
		rconn.connect()
		return rconn


	def __getslots(self):
		cmd = [
			"/usr/bin/rbd",
			"ls",
			"--pool", self.config['pool'],
			"--format","json"
		]
		#print cmd
		status,output = self.__run_command(cmd)
		output = json.loads(output)
		return output

	def __getmapped_slot(self, slot):
		mapped_device = None
		mapped = self.__getmapped()
		#print mapped
		for mapid,mapped in mapped.items():
			#print mapid
			#print mapped['device']
			m_int = int(mapped['name'].replace(self.config['prefix'],''))
			#print m_int, slot
			if m_int == slot:
				#print "Already mapped"
				
				return mapped
		return None

	def __map(self, slot):
		cmd = [
			'rbd',
			'map',
			'%s%.5d'%(self.config['prefix'], slot),
			'--pool', self.config['pool']
			]
		status,output = self.__run_command(cmd)
		if status != 0:
			print output
			raise RuntimeError("Error status %s from '%s'"%(status, " ".join(cmd)))

	def __unmap(self, device):
		cmd = [
			'rbd',
			'unmap',
			device
			]
		status,output = self.__run_command(cmd)
		if status != 0:
			print output
			raise RuntimeError("Error status %s from '%s'"%(status, " ".join(cmd)))

	def labelnew(self, slot, path, driveindex):
		rconn = self.__open_rados()
		print "Pool: %s"%(self.config['pool'])
		ioctx = rconn.open_ioctx(str(self.config['pool']))
		try:
			counter = ioctx.read('counter')
		except rados.ObjectNotFound:
			counter = "0"
			#ioctx.write_full('counter','0')

		counter = str( int(counter) + 1 )
		ioctx.write_full('counter',str(counter))

		label = "%s%.5d"%(self.config['prefix'], int(counter))
		print "New label: %s"%(label)

		rbdconn = rbd.RBD()
		try:
			rbdconn.create(ioctx, str(label), int(self.config['imagesize'])*1024*1024)
		except rbd.ImageExists:
			print "Image '%s' already exists! Skipping!"%(label)

		# Map image to rbd
		self.__map(int(counter))

		# Format image
		if self.config['imageformat'] == 'ext4':
			cmd = [
				'mke2fs', '-j', '-O','extents',
				'-L', str(label),
				'/dev/rbd/%s/%s'%(self.config['pool'], label)
			]
		else:
			raise ValueError("Unknown image format '%s'"%(self.config['imageformat']))

		status,output = self.__run_command(cmd)