#! /usr/bin/python

import sys
import commands
import os
import re

#We're going to do nasty shit, need root
if not os.geteuid() == 0:
    sys.exit('Error: pininit must be run as root. Please run sudo pininit')

#Load up the config file
if not os.path.exists("pin.conf"):
    sys.exit("Error: Cannot find pin.conf")

#Pin conf (should) contains valid python setting "default" and "pinned" groups up
default = []
pinned  = []
pinconf = "".join(open("pin.conf").readlines())
exec(pinconf)

if default == []:
    sys.exit("Error: Default CPU list is empty. Check that pin.conf has the line default = [ <CPU NUMS> ] in it")

if pinned == []:
    sys.exit("Error: Pinned CPU list is empty. Check that pin.conf has the line pinned = [ <CPUS NUMs> ] in it")

#first build the hex mask
hexcpumask = map(lambda x: 1 << x , default)
hexcpumask = reduce(lambda x,y: x + y, hexcpumask)
print hexcpumask

#Get the current list of processes 
out = commands.getoutput('ps -A')
lines = out.split("\n")
findpid = re.compile("\s*(\d*)\s.*")
pids = []
for line in lines:
    result = findpid.search(line)
    if(result):
        if(result.group(1)):
            pids.append(result.group(1))

#Move all pids on the defualt CPU set
for pid in pids:
    cmd = "taskset -p %x %s" % (hexcpumask,pid)
    print "Running \"" + cmd + "\" ..."
    out = commands.getoutput(cmd)


#Get all the IRQs
procirqs = open("/proc/interrupts").readlines()
findirq = re.compile("\s*(\d*):.*")
irqs = []
for line in procirqs:
    result = findirq.search(line)
    if(result):
        if(result.group(1)):
            irqs.append(result.group(1))
irqs = []

#Move all irqs on the defualt CPU set
for irq in irqs:
    cmd = "echo %x > /proc/irq/%s/smp_affinity" % (hexcpumask,irq)
    print "Running \"" + cmd + "\" ..."
    out = commands.getoutput(cmd)

#Set the defualt irq affinit to the default cpu set
cmd = "echo %x > /proc/irq/default_smp_affinity" % (hexcpumask)
#print "Running \"" + cmd + "\" ..."
#out = commands.getoutput(cmd)



#Put this somewhere global so that pin can use it
pinned = open("/tmp/pinned","w+")
pinned.write("inuse=[" + ",".join(map(lambda x: str(x), default)) + "] \n")
pinned.close()



