#!/usr/bin/python
#
# Copyright (C) 2012 STEC, Inc. All rights not specifically granted
# under a license included herein are reserved
# Wrote a python based CLI for admistration of Enhanceio Driver
# Sanoj Unnikrishnan <sunnikrishnan@stec-inc.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; under version 2 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.


from ctypes import *
from fcntl import *
from argparse import ArgumentParser
import sys,struct
import subprocess
import fcntl
import os
import re

libc = CDLL('libc.so.6')

#TBD : Change ioctl numbers to comply with linux kernel convention
EIODEV = '/dev/eiodev'
EIO_IOC_CREATE = 1104168192 
EIO_IOC_DELETE = 1104168193
EIO_IOC_ENABLE = 1104168194
EIO_IOC_DISABLE = 1104168195
EIO_IOC_EDIT = 1104168196
EIO_IOC_NCACHES = 1104168197
EIO_IOC_CACHE_LIST = 1104168198
EIO_IOC_SSD_ADD = 1104168199
EIO_IOC_SSD_REMOVE = 1104168200
EIO_IOC_SRC_ADD = 1104168201
EIO_IOC_SRC_REMOVE = 1104168202
IOC_BLKGETSIZE64 = 0x80081272
IOC_SECTSIZE = 0x1268
SUCCESS=0
FAILURE=3

FIOCMD="fio --rw=randrw --verify=meta --verify_pattern=0x3982 --ioengine=libaio \
	--iodepth=16 --direct=1 --size=1G --bsrange=512-128K --numjobs=10 \
	--group_reporting --time_based --runtime=180 --name=job1 --filename="

udev_template = """
ACTION!="add|change", GOTO="EIO_EOF"
SUBSYSTEM!="block", GOTO="EIO_EOF"

<cache_match_expr>, GOTO="EIO_CACHE"
<source_match_expr>, GOTO="EIO_SOURCE"

# If none of the rules above matched then it isn't an \
EnhanceIO device so ignore it.
GOTO="EIO_EOF"

#=================== EIO_CACHE =======================

# If we just found the cache device and the source already \
exists then we can setup

LABEL="EIO_CACHE"
TEST!="/dev/enhanceio/<cache_name>", \
PROGRAM="/bin/mkdir -p /dev/enhanceio/<cache_name>"
PROGRAM="/bin/sh -c 'echo $kernel > /dev/enhanceio/<cache_name>/.ssd_name'"
TEST=="/proc/enhanceio/<cache_name>", \
RUN+="/sbin/eio_cli notify -a add -s /dev/$kernel -c <cache_name>", \
GOTO="EIO_EOF"

TEST=="/dev/enhanceio/<cache_name>/.disk_name", GOTO="EIO_SETUP"

GOTO="EIO_EOF"

#=================== EIO_SOURCE =======================

# If we just found the source device and the cache already \
exists then we can setup

LABEL="EIO_SOURCE"
TEST=="/dev/enhanceio/<cache_name>/.eio_delete", GOTO="EIO_SOURCE_DELETE"
TEST!="/dev/enhanceio/<cache_name>", PROGRAM="/bin/mkdir -p \
/dev/enhanceio/<cache_name>"
PROGRAM="/bin/sh -c 'echo $kernel > /dev/enhanceio/<cache_name>/.disk_name'"

PROGRAM="/bin/sh -c 'echo $major $minor > \
/dev/enhanceio/<cache_name>/.disk_mjr_mnr'"

PROGRAM="/bin/sh -c 'echo $links > \
/dev/enhanceio/<cache_name>/.srclinks'"

TEST=="/proc/enhanceio/<cache_name>",\
PROGRAM="/bin/grep 'state\s*failed' /proc/enhanceio/<cache_name>/config", \
PROGRAM="/bin/grep 'no_source_dev\s*1' /proc/enhanceio/<cache_name>/errors", \
RUN+="/sbin/eio_cli notify -a add -d /dev/$kernel -c <cache_name>", NAME="", \
GOTO="EIO_EOF"

TEST!="/proc/enhanceio/<cache_name>", \
TEST!="/dev/enhanceio/<cache_name>/.eio_delete", ACTION!="change", NAME=""
	
TEST=="/dev/enhanceio/<cache_name>/.ssd_name", GOTO="EIO_SETUP"

GOTO="EIO_EOF"

#=================== EIO_SOURCE_DELETE =======================

LABEL="EIO_SOURCE_DELETE"
PROGRAM="/bin/sh -c 'cat /dev/enhanceio/<cache_name>/.disk_name'", \
ENV{disk_name}="%c"
TEST=="/dev/enhanceio/<cache_name>/.disk_name", ENV{disk_name}=="?*", \
NAME="$env{disk_name}"
PROGRAM="/bin/unlink /dev/enhanceio/<cache_name>/.disk_name"
PROGRAM="/bin/unlink /dev/enhanceio/<cache_name>/.disk_mjr_mnr"
GOTO="EIO_EOF"
#=================== EIO_SETUP =======================

LABEL="EIO_SETUP"
PROGRAM="/bin/sh -c 'cat /dev/enhanceio/<cache_name>/.ssd_name'", \
ENV{ssd_name}="%c"
PROGRAM="/bin/sh -c 'cat /dev/enhanceio/<cache_name>/.disk_name'", \
ENV{disk_name}="%c"

TEST!="/proc/enhanceio/<cache_name>", \
TEST!="/dev/enhanceio/<cache_name>/.skip_enable", \
RUN+="/bin/sh -c '/bin/mknod /dev/$env{disk_name} b \
`/bin/cat /dev/enhanceio/<cache_name>/.disk_mjr_mnr`'"

TEST!="/proc/enhanceio/<cache_name>", \
RUN+="/bin/sh -c ' for i in `cat /dev/enhanceio/<cache_name>/.srclinks`; \
do rm -f /dev/$$i; ln -f -s /dev/$env{disk_name} /dev/$$i; done'"


TEST!="/proc/enhanceio/<cache_name>", RUN+="/sbin/eio_cli \
enable -d /dev/$env{disk_name} -s /dev/$env{ssd_name} -m <mode> \
-b <block_size> -p <policy> -c <cache_name>"

LABEL="EIO_EOF"
"""

def make_udev_match_expr(dev_path, cache_name):
	dict_udev = {}
	status = run_cmd("udevadm info --query=property --name=" + dev_path)
	for line in status.output.split('\n'):
		if line is not '':
			temp = line.split("=", 1)
			dict_udev[temp[0].strip()] = temp[1].strip()

	# DM devices
	if dict_udev["DEVTYPE"] == "disk" and "DM_UUID" in dict_udev:
		match_expr='ENV{DM_UUID}=="'+ dict_udev["DM_UUID"] + \
			 '", ENV{DEVTYPE}=="' + dict_udev["DEVTYPE"] +'"'
		return	match_expr

	# MD devices
	if "MD_UUID" in dict_udev:
		if dict_udev["DEVTYPE"] == "disk" and \
		"MD_DEV_UUID" not in dict_udev: 
			match_expr='ENV{MD_UUID}=="' + dict_udev["MD_UUID"] + \
			'", ENV{DEVTYPE}=="' + dict_udev["DEVTYPE"] + '"'
			return match_expr

		elif dict_udev["DEVTYPE"] == "partition":
			try:
				with open("/sys" + dict_udev["DEVPATH"] + "/partition") as f:
					partition_num = f.read().strip()
					match_expr='ENV{MD_UUID}=="' + \
					'", ENV{MD_DEV_UUID}=="' + \
					'", ATTR{partition}=="' + partition_num + '"'
					return match_expr	
			except IOError as e:
				pass

	# Partition
	if dict_udev["DEVTYPE"] == "partition" and "ID_SERIAL" in dict_udev:
		try:
			with open("/sys" + dict_udev["DEVPATH"] + "/partition") as f:
				partition_num = f.read().strip()
				match_expr='ENV{ID_SERIAL}=="' + \
				dict_udev["ID_SERIAL"] + \
				'", ATTR{partition}=="' + partition_num + '"'
				return match_expr
		except IOError as e:
			pass


	# Disk
	if dict_udev["DEVTYPE"] == "disk" and "ID_SERIAL" in dict_udev :
		match_expr='ENV{ID_SERIAL}=="'+ dict_udev["ID_SERIAL"] + \
		'", ENV{DEVTYPE}=="'+ dict_udev["DEVTYPE"] +'"'
		return match_expr

	# Partition or disk w/ filesystem
	if "ID_FS_UUID" in dict_udev:
		match_expr= 'ENV{DM_NAME}!="' + cache_name + \
		'", ENV{ID_FS_UUID}=="' + dict_udev["ID_FS_UUID"] + '"'
		return match_expr


def run_cmd(cmd):
	#Utility function that runs a command
	process = subprocess.Popen(cmd, stdout=subprocess.PIPE,\
				stderr=subprocess.STDOUT,shell=True)
	output = process.stdout.read()
	ret = process.wait()
	status = Status(output,ret)
	return status



def get_caches_list():
	
	#Utility function that obtains cache list 
	cache_list = [f for f in os.listdir('/proc/enhanceio/')]
	for name in cache_list:
		if name == "version":
			cache_list.remove(name)
	return cache_list

def sanity(hdd, ssd):
	# Performs a very basic regression of operations			
				
	modes = {3:"Write Through", 1:"Write Back", 2:"Read Only",0:"N/A"}
	policies = {3:"rand", 1:"fifo", 2:"lru", 0:"N/A"}		
	blksizes = {"4096":4096, "2048":2048, "8192":8192,"":0}	
	for mode in ["wb","wt","ro"]:
		for policy in ["rand","fifo","lru"]:
			for blksize in ["4096","2048","8192"]:
				cache = Cache_rec(name = "test_cache", src_name = hdd,\
						ssd_name = ssd, policy = policy, mode = mode,\
						blksize = blksize)
				cache.create()
				run_cmd(FIOCMD + hdd)
				cache.delete()
	

# Class that represents cache. Also used to pass ioctl to driver 
class Cache_rec(Structure):
	_fields_ = [
	("name", c_char * 32),
	("src_name", c_char * 128),
	("ssd_name", c_char * 128),
	("ssd_uuid", c_char * 128),
        ("src_size", c_ulonglong),	
	("ssd_size", c_ulonglong),
        ("src_sector_size", c_uint),
        ("ssd_sector_size", c_uint),
	("flags", c_uint),
        ("policy", c_byte),
        ("mode", c_byte),
        ("persistence", c_byte),
        ("cold_boot", c_byte),
        ("blksize", c_ulonglong),
        ("assoc", c_ulonglong)
	]
	def __init__(self, name, src_name="", ssd_name="", src_size=0,\
		     ssd_size=0, src_sector_size=0, ssd_sector_size=0,\
		     flags=0, policy="", mode="", persistence=0, cold_boot="",\
		     blksize="", assoc=""): 
	
		modes = {"wt":3,"wb":1,"ro":2,"":0}
		policies = {"rand":3,"fifo":1, "lru":2,"":0}
		blksizes = {"4096":4096, "2048":2048, "8192":8192,"":0}	
		associativity = {2048:128, 4096:256, 8192:512,0:0}
		
		self.name = name
		self.src_name =src_name
		self.ssd_name = ssd_name
		self.src_size = src_size	
		self.src_sector_size = src_sector_size
		self.ssd_size = ssd_size
		self.ssd_sector_size = ssd_sector_size
		self.flags = flags
		self.policy = policies[policy] 
		self.mode = modes[mode]
		self.persistence = persistence
		self.blksize = blksizes[blksize]
		self.assoc = associativity[self.blksize]		 

	
	def print_info(self):
	
		# Display Cache info 
		modes = {3:"Write Through", 1:"Write Back", 2:"Read Only",0:"N/A"}
		policies = {3:"rand", 1:"fifo", 2:"lru", 0:"N/A"}

		print "Cache Name       : " + self.name 
		print "Source Device    : " + self.src_name 
		print "SSD Device       : " + self.ssd_name
		print "Policy           : " + policies[self.policy] 
		print "Mode             : " + modes[self.mode]
		print "Block Size       : " + str(self.blksize)	
		print "Associativity    : " + str(self.assoc)	

		if os.path.exists("/proc/enhanceio/" + self.name):
			cmd = "cat /proc/enhanceio/" + self.name + "/config" + " | grep state"
			status = run_cmd(cmd)
			print "State            : " + status.output.split()[1]
	
		pass


	def do_eio_ioctl(self,IOC_TYPE):
		#send ioctl to driver
		fd = os.open(EIODEV, os.O_RDWR, 0400)
		selfaddress = c_ulong(addressof(self))
		try:
			if libc.ioctl(fd, IOC_TYPE, selfaddress) == SUCCESS:
				return SUCCESS
		except Exception as e:
			print e
		return FAILURE

		
	def clean(self):
		#do sysctl corresponding to clean
		cmd = "/sbin/sysctl dev.enhanceio." + self.name + ".do_clean=1"
		print cmd
		run_cmd(cmd)	
		pass

	
	def get_cache_info(self):
		#function to extract information from /proc/enhanceio
		status = Status()
		
		if os.path.exists("/proc/enhanceio/" + self.name):

			associativity = {2048:128, 4096:256, 8192:512,0:0}
				
			cmd = "cat /proc/enhanceio/" + self.name + "/config" + " | grep src_name" 
			status = run_cmd(cmd)
			self.src_name =  status.output.split()[1]

			cmd = "cat /proc/enhanceio/" + self.name + "/config" + " | grep ssd_name"
			status = run_cmd(cmd)
			self.ssd_name = status.output.split()[1]
	
			cmd = "cat /proc/enhanceio/" + self.name + "/config" + " | grep mode"
			status = run_cmd(cmd)
			self.mode = int(status.output.split()[1])

			cmd = "cat /proc/enhanceio/" + self.name + "/config" + " | grep eviction"
			status = run_cmd(cmd)
			self.policy = int(status.output.split()[1])

			cmd = "cat /proc/enhanceio/" + self.name + "/config" + " | grep block_size"
			status = run_cmd(cmd)
			self.blksize = int(status.output.split()[1])

			self.assoc = associativity[self.blksize]
			
			return SUCCESS

		return FAILURE


	def create(self):		 

		self.print_info()	
		if self.do_eio_ioctl(EIO_IOC_CREATE) == SUCCESS:
			self.create_rules()
			print 'Cache created successfully'
			return SUCCESS
		else: 	
			print 'Cache creation failed ' + \
			'(dmesg can provide you more info)'
	
		return FAILURE

		
	def edit(self):
		mode = self.mode
		policy = self.policy
		if self.get_cache_info() == SUCCESS:
			if mode != 0:
				self.mode = mode
			if policy != 0:
				self.policy = policy
			
			if self.do_eio_ioctl(EIO_IOC_EDIT) == SUCCESS:
				self.create_rules()
				print 'Cache edited Successfully'
				return SUCCESS
			else:
				print 'Edit cache failed ' + \
				'(dmesg can provide you more info)'
		else:
			print 'Requested cache not found'
		
		return FAILURE


	def create_rules(self):
		
		source_match_expr = make_udev_match_expr(self.src_name, self.name)
		print source_match_expr
		cache_match_expr = make_udev_match_expr(self.ssd_name, self.name)
		print cache_match_expr
		modes = {3:"wt", 1:"wb", 2:"ro",0:"N/A"}
		policies = {3:"rand", 1:"fifo", 2:"lru", 0:"N/A"}
	
		try: 	
			udev_rule = udev_template.replace("<cache_name>",\
				    self.name).replace("<source_match_expr>",\
				    source_match_expr).replace("<cache_match_expr>",\
				    cache_match_expr).replace("<mode>",\
				    modes[self.mode]).replace("<block_size>",\
				    str(self.blksize)).replace("<policy>",\
				    policies[self.policy])	

			# write rule file
			rule_file_path = "/etc/udev/rules.d/94-enhanceio-" + \
					self.name + ".rules" 
			rule_file = open(rule_file_path, "w")
			rule_file.write(udev_rule)
			return SUCCESS

		except Exception as e:
			print "Creation of udev rules file failed\n" 
			print e
 			return FAILURE

	def delete(self):

		if self.get_cache_info() == FAILURE:
			print "No cache exists with the name " + self.name
			return FAILURE

		run_cmd("/bin/touch /dev/enhanceio/" + self.name + "/.eio_delete")

		if self.do_eio_ioctl(EIO_IOC_DELETE) == SUCCESS:
			if self.delete_rules() == SUCCESS:
				return SUCCESS
		else: 	
			print 'Cache deletion failed (dmesg can provide you more info)'
		
		run_cmd("/bin/rm -f /dev/enhanceio/" + self.name + "/.eio_delete")
		
		return FAILURE

						
	def delete_rules(self):
				
		rule_file_path = "/etc/udev/rules.d/94-enhanceio-" + self.name + ".rules" 
		print "Removing file" + rule_file_path
		try:	
			os.remove(rule_file_path)
			print "Done removing rule file"
			return SUCCESS
		except Exception as e:
			print "Unable to delete rule file" + \
			"(please remove it manually)"
			print e
		return FAILURE
	
class Status:
	output = ""
	ret = 0

	def __init__(self, outstr="", outret=0):
		self.output = outstr
		self.ret = outret			
		pass

def create_parser():

	mainparser = ArgumentParser()
	parser = mainparser.add_subparsers()		

	#delete
	parser_delete = parser.add_parser('delete', help='used to delete cache')
	parser_delete.add_argument("-c", action="store", dest= "cache",required=True)

	#edit
	parser_edit = parser.add_parser('edit', help="used to edit \
					cache policy or mode or both")
	parser_edit.add_argument("-c", action="store",  dest="cache",required=True)
	parser_edit.add_argument("-m", action="store", dest="mode", \
			choices=["wb","wt","ro"], help="cache mode",default="")
	parser_edit.add_argument("-p", action="store", dest="policy", \
				choices=["rand","fifo","lru"], help="cache \
				replacement policy",default="") 
	
	#info
	parser_info = parser.add_parser('info', help='displays information \
					about currently create caches')
	
	#clean
	parser_clean = parser.add_parser('clean', help='clean the dirty blocks \
				in the cache (Applicable only to writeback caches)')
	parser_clean.add_argument("-c", action="store", dest="cache",required=True)
	
	#create
	parser_create = parser.add_parser('create', help="create")
	parser_create.add_argument("-d", action="store", dest="hdd",\
				required=True, help="name of the source device")
	parser_create.add_argument("-s", action="store", dest="ssd",\
				required=True, help="name of the ssd device")
	parser_create.add_argument("-p", action="store", dest="policy",\
				   choices=["rand","fifo","lru"],\
				   help="cache replacement policy",default="lru")
	parser_create.add_argument("-m", action="store", dest="mode",\
				   choices=["wb","wt","ro"],\
				   help="cache mode",default="wt")
	parser_create.add_argument("-b", action="store", dest="blksize",\
				   choices=["2048","4096","8192"],\
				   default="4096" ,help="block size for cache")
	parser_create.add_argument("-c", action="store", dest="cache", required=True)
	
	#enable
	parser_enable = parser.add_parser('enable', help='used to enable cache')
	parser_enable.add_argument("-d", action="store", dest="hdd",\
				   required=True, help="name of the source device")
	parser_enable.add_argument("-s", action="store", dest="ssd",\
				   required=True, help="name of the ssd device")
	parser_enable.add_argument("-p", action="store", dest="policy",
				   choices=["rand","fifo","lru"],\
				   help="cache replacement policy",default="lru")
	parser_enable.add_argument("-m", action="store", dest="mode",\
				   choices=["wb","wt","ro"],\
				   help="cache mode",default="wt")
	parser_enable.add_argument("-b", action="store", dest="blksize",\
				   choices=["2048","4096","8192"],\
				   default="4096" ,help="block size for cache")
	parser_enable.add_argument("-c", action="store", dest="cache", required=True)

	#notify
	parser_notify = parser.add_parser('notify')
	parser_notify.add_argument("-s", action="store", dest="ssd",\
				   help="name of the ssd device")
	parser_notify.add_argument("-a", action="store", dest="action",\
				    help="add/remove/reboot")
	parser_notify.add_argument("-d", action="store", dest="hdd",\
				    help="name of the source device")
	parser_notify.add_argument("-c", action="store", dest="cache", required=True)

	#sanity
	parser_sanity = parser.add_parser('sanity')
	parser_sanity.add_argument("-s", action="store", dest="ssd",\
			          required=True, help="name of the ssd device")
	parser_sanity.add_argument("-d", action="store", dest="hdd",\
				  required=True, help="name of the source device")

	#disable
	parser_disable = parser.add_parser('disable', help='used to disable cache')
	parser_disable.add_argument("-c", action="store", dest="cache", required=True)

	return mainparser

def main():
	 	
	mainparser = create_parser()
	args = mainparser.parse_args()

	run_cmd("/sbin/modprobe enhanceio")
	run_cmd("/sbin/modprobe enhanceio_fifo")
	run_cmd("/sbin/modprobe enhanceio_lru")
	run_cmd("/sbin/modprobe enhanceio_rand")

	f = open("/run/eio.lock", "w")
	fcntl.flock(f, fcntl.LOCK_EX)

	if sys.argv[1] == "create":

		if re.match('^[\w]+$', args.cache) is None:
			print "Cache name can contain only alphanumeric" + \
			" characters and underscore ('_')"
			return FAILURE

		cache = Cache_rec(name = args.cache, src_name = args.hdd,\
				ssd_name = args.ssd, policy = args.policy,\
				mode = args.mode, blksize = args.blksize)
		return cache.create()

	elif sys.argv[1] == "info":
		cache_list = get_caches_list()
		if not cache_list:		
			print "No caches Found"
		else:
			for cache_name in cache_list:
				cache = Cache_rec(name = cache_name)
				cache.get_cache_info()
				cache.print_info()

			print "\nFor more information look at " + \
			      "/proc/enhanceio/<cache_name>/config"
		return SUCCESS

	elif sys.argv[1] == "edit":
		cache = Cache_rec(name = args.cache, policy = args.policy, mode = args.mode)
		return cache.edit()

	elif sys.argv[1] == "delete":
		cache = Cache_rec(name = args.cache)
		return cache.delete()

	elif sys.argv[1] == "clean":
		cache = Cache_rec(name = args.cache)
		return cache.clean()

	elif sys.argv[1] == "enable":
		# This command will be fired by udev rule on SSD/Source addition
		cache = Cache_rec(name = args.cache, src_name = args.hdd,\
				ssd_name = args.ssd, policy = args.policy,\
				mode = args.mode, blksize = args.blksize, persistence = 1)

		cache.do_eio_ioctl(EIO_IOC_ENABLE)

	elif sys.argv[1] == "disable":
		cache = Cache_rec(name = args.cache)
		cache.do_eio_ioctl(EIO_IOC_DISABLE)

	elif sys.argv[1] == "notify":
		# This command will be fired by udev rule on SSD/Source addition
		if args.action == "reboot":
			cache = Cache_rec(name = "dummy")
			cache.do_eio_ioctl(EIO_IOC_ENABLE)

		elif args.action == "add": 
			if args.ssd:
				cache = Cache_rec(name = args.cache,\
						ssd_name = args.ssd, persistence = 1)
				cache.do_eio_ioctl(EIO_IOC_SSD_ADD)
				
			elif args.hdd:
				cache = Cache_rec(name = args.cache,\
						src_name = args.hdd, persistence = 1)
				cache.do_eio_ioctl(EIO_IOC_SRC_ADD)

		elif args.action == "remove":
			if args.ssd:
				cache = Cache_rec(name = args.cache, ssd_name = args.ssd,\
						persistence = 1)
				cache.do_eio_ioctl(EIO_IOC_SSD_REMOVE)
			elif args.hdd:
				cache = Cache_rec(name = args.cache, src_name = args.hdd,\
						persistence = 1)
				cache.do_eio_ioctl(EIO_IOC_SRC_REMOVE)

		pass

	elif sys.argv[1] == "sanity":
		# Performs a basic sanity check
		sanity(args.hdd, args.ssd)
		return SUCCESS

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

