#!/usr/bin/env python
#
#
#==========================================================================
#  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; either version 2 of the License, or    
#  (at your option) any later version.                                  
#==========================================================================
#
# $Id: limitcpu.py,v 1.1.1.1 2009-10-18 12:19:02 main Exp $
# Eri Ramos Bastos - bastos.eri () gmail.com
#
from subprocess import Popen, PIPE, STDOUT
import os
import sys


### Define Processes and their limits here.
# Should be "process_name":limit
# Where: process name is a regular expression used to find the pids
#        limit is a % value of each available core.
# Therefore: "VIN":20 will allow each VIN process to use up to 20% of a CPU core.
#            "wsh":180 will allow each wsh process to use up to 100% of a CPU core + 80% of a second core

process_n_limits={"firefox":50, "VirtualBox":20}


################# Don't change from here to bottom #########################


class LimitCpu(Exception):
    def __init__(self,processname,limit):
        self.processname = processname
        self.limit = limit


    def GetProcesses(self):

        """ 
            Get all processes that match the REGEX provided 
            and returns a dictionary with process ids and current CPU usage in %.
        """
        # Variables
        self.processes = {}

        command = 'top -b -n1 -c|grep %s|grep -v grep' % self.processname
        running_processes = Popen(command, shell=True,stdout=PIPE,stderr=PIPE)
        for line in running_processes.stdout.readlines():
            self.processes[line.split()[0]] = line.split()[8]
        return self.processes


    def convertStr(self,s):
        """Convert string to either int or float."""
        try:
            ret = int(s)
        except ValueError:
            ret = float(s)
        return ret


    def CheckIfManaged(self):

        """
            Check if a list of PIDs are currently being managed by cpulimit.
               If PID CPU usage is under limit, skip it.
            Returns 3 lists: First for unmaneged processes, Second for managed ones
               and Third for under limit.
        """
        # Variables
        self.unmanaged = []
        self.managed = []
        self.under_limit = []

        command = "ps auxww|grep 'cpulimit -p %s -l'|grep -v grep"

        processes = self.GetProcesses()
        for pid in processes:
            if self.convertStr(processes[pid]) < self.limit:
                self.under_limit.append(pid)
            else:
                running_processes = Popen(command % pid, shell=True,stdout=PIPE,stderr=PIPE).stdout.readlines()
                if running_processes:
                    self.managed.append(pid)
                else:
                    self.unmanaged.append(pid)
        
        return self.unmanaged, self.managed, self.under_limit


    def ApplyLimit(self):
        """
            Applies limit to a list of PIDs
        """

        command = "/usr/bin/cpulimit -p %s -l %s -z &"
        for pid in self.CheckIfManaged()[0]:
            command_result = os.system(command % (pid, self.limit))

for process in process_n_limits:
    managed_pid = LimitCpu(process,process_n_limits[process])
    managed_pid.ApplyLimit()