#!/usr/bin/env python

#  Disker1 -- a prreliminary to DrakOPT--an optimization program
#  for IDE HDDs.
#  For testing purposes only--released as soon as coded
#  Copyright 2001 by civileme@mandrakesoft All rights reserved
#  Licensed under the GNU GPL  This code may NOT be reissued
#  as shareware or freeware, only under the terms of the current
#  GNU General Public license.  This is NOT Public domain open
#  for exploitation.
#
#  NO WARRANTY WHATSOEVER.  If your computer throws a lightning
#  bolt at you or polymorphs you into a bandersnatch, you better
#  run from the Jabberwock.



import user
import string
import imp
import types
import os
import sys
import re
global disklist, dmadic, dmardic,  lookah, multah, speedah, speedax, dmarrdic, seeker, finder
lookah = {0:0,1:2,2:4,3:8}
multah ={0:0,1:2,2:8,3:16}
speedah=[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]
def scorespeed(xindex,lindex,mindex,speed):
    global  speedax
    speedax[xindex*16+lindex*4+mindex]=speed
    return
def bespeed():
    global lookah, multah, speedax
    maxim=0.0
    maxidx=0
    j=0
    for i in speedax:
        if i >= maxim:
            maxim=i
            maxid=j
        j=j+1
    xid=int(maxid/16)
    aid=maxid-xid*16
    lid=int(aid/4)
    mid=aid-lid*4
    lset=lookah[lid]
    mset=multah[mid]
    return lset, mset, xid, maxim
def parmset(s,uset,cset,dset,mset,aset,rateset):
    umask=' -u'
    c32=' -c'
    multcnt=' -m'
    readahead=' -a'
    rate=' -X'
    mks='hdparm'+umask+repr(uset)+c32+repr(cset)+' -d1'+multcnt+repr(mset)+readahead+repr(aset)+rate+repr(rateset)+' /dev/'+s
    os.system('sync') # Get the settings from previous step saved if any
    os.system(mks+' > /etc/idedrives/.drives/results')
    os.system(mks+' > /etc/idedrives/.drives/results') #second one is the write test--if we bomb here, safe boot
    return mks

def speedtest(s):
    global seeker, finder
    finder=[]
    seeker=os.system('hdparm -t /dev/'+s+' > /etc/idedrives/.drives/results')
    if seeker !=0:
        inputbad=open('/etc/idedrives/.drives/results','r')
        finder=inputbad.readlines()
        print
        print 'errors encountered--we will not proceed with program'
        print
        return -1.0
    os.system('hdparm -t /dev/'+s+' > /etc/idedrives/.drives/results')
    os.system('hdparm -t /dev/'+s+' > /etc/idedrives/.drives/results') # Test reading three times
    os.system('sync')
    inputres=open('/etc/idedrives/.drives/results','r')
    resultx=inputres.readlines()
    result=resultx[2]
    fresult=string.split(result)
    nhdparmumber=float(fresult[10])
    return nhdparmumber



def erhandler(s, scompanion):
    global seeker, finder
    # if we get here, we have a problem/error message from hdparm -t which is very bad
    s1='^'+s
    s2='^'+scompanion
    badmessage1=0
    badmessage2=0
    for g in finder:
        if re.search(s1,g):
            badmessage1=1
        if re.search(s2,g):
            badmessage2=2
    glory=badmessage1+badmessage2
    if glory >2:
        print
        print 'SERIOUS PROBLEM WITH HARDWARE'
        print 'Crosstalk has been detected'
        print 'Split the drives '+s+' from '+scompanion
        print 'So they no longer share the same channel'
    elif glory >= 1:
        print
        print 'This is not good--a hardware problem'
        print 'exists with cable or drive '+s
        print 'or with the associated ide channel'
        print 'we have backed off the setting'
        print
    else:
        pass
    return glory
    

    
def salt(xcess):
   
    inputbuf=open('/etc/rc.d/init.d/drivetune','r')
    george=inputbuf.readlines()
    inputbuf.close()
    outputbuf=open('/etc/rc.d/init.d/drivetune','w')
    for gorgon in george:
        if string.find(gorgon,xcess[-6:]) < 0:
            outputbuf.write(gorgon)
    outputbuf.write(xcess+'>> /dev/null\n')
    outputbuf.close()
    return

dmadic= { 'DMA': 33, 'DMA2': 34, 'UDMA': 65, 'UDMA2': 66, 'UDMA3': 67, 'UDMA4': 68, 'UDMA5': 69 }
dmardic=dmadic.keys()
dmardic.sort()
dmarrdic={33:0,34:1,65:2,66:3,67:4,68:5,69:6}

disklist=['s','hda','hdb','hdc','hdd','hde','hdf','hdg','hdh','hdi','hdj','hdk','hdl','hdm','hdn','hdo','hdp']

    
print "                   Welcome to DrakOPT"
os.system('uname -a > /tmp/.unameXJZ')
inputu=open('/tmp/.unameXJZ','r')
unamel=inputu.readline()
unames=string.split(unamel)
print unames
inputu.close()
os.system('rm /tmp/.unameXJZ -f')
uname=unames[2]

if uname[2:3]!='4':
    print
    print 'There is no way this program can run.  The hdparm instruction will '
    print 'have no capabilities line.  Kernel 2.4 must be in use'
    print 'Bye'
    sys.exit(0)
else:
    print
    print 'Kernel 2.4 verified'
    print

print
print "      This module will take some time to run; 50 "
print "      MINUTES is not uncommon?"
print
print "In rare circumstances, it could cause your machine"
print "to freeze up.  If this happens, reboot and your settings"
print "will be at a safe level."
print
print "  ----DON'T RUN THIS WITH Other Programs----"
print " --In fact, mouse movements can cause inaccuracies so "
print " Just watch the show and have your favorite beverage---"
print
print "Disker uses a sweet-spot seeking algorithm to set each of"
print "your disk drives to a safe, but efficient level of "
print "operation.  This is often a lower setting than the "
print "manufacturer claims for his drive."
print
print "You may override any of the Disker settings by entering"
print "your own hdparm settings in /etc/drives.local"
print 
print "OPTIONS:"
print "         (A)--Abort"
print "         (N)--New (or you changed disks around)"
print "         (C)--Continue from a previous run"
print "Your Choice (N/c/a)"
answer=raw_input('')
search=0
if len(answer) < 1:
    answer='N'
if answer[0]=='N' or answer[0] =='n':
            
    os.system(' echo "#!/bin/sh" > /etc/rc.d/init.d/drivetune')
    os.system(' echo "# Drive tuning hdparms loaded by DrakOPT" >> /etc/rc.d/init.d/drivetune')
    r='echo -n "Implementing IDE Optimisations"'
    os.system('echo '+r+' >> /etc/rc.d/init.d/drivetune')
    os.system(' chmod 0744 /etc/rc.d/init.d/drivetune ')
    os.system('sync')
    os.system('ln -fs /etc.rc.d/init.d/drivetune.sh /etc/rc.d/rc3.d/S98drivetune')
    os.system('ln -fs /etc.rc.d/init.d/drivetune.sh /etc/rc.d/rc5.d/S98drivetune')
    os.system('rm -r /etc/idedrives/.drives -f')
    os.system('[ -d /etc/idedrives/.drives ] || mkdir /etc/idedrives/.drives  ')
    os.system('echo "y" > /etc/idedrives/.drives/s')    
elif answer[:1]=='c' or answer[:1]=='C':
    if os.system('[ -d /etc/idedrives/.drives ] '):
        print
        print "Ooops--I think you meant new.  Try this program again"
        print
        sys.exit(0)
    pass
else:
    sys.exit(0)

input=open('/etc/idedrives/.drives/s','r')
input=open('/etc/idedrives/.drivedb','r')
drivedb=input.readlines()
beta=[]
count=0
for hj in drivedb:
    alpha=string.split(hj)
    beta=beta+[string.replace(string.upper(hj[1]),' ','')]
    count=count+1
    
for i in range(16):
    j=i+1
    s=disklist[j]
    kk=int(j/2)
    if j-2*kk == 0:
        companion=j-1
    else:
        companion=j+1
    scompanion=disklist[companion]
    
    if os.system('[ -e /etc/idedrives/.drives/'+s+' ] ') == 0:  # already worked?  
        continue                             # Skip it
    os.system('touch /etc/idedrives/.drives/'+s)  # close the door for next time
    hdmedia=''
    # if we reach here, we want to configure a drive, but it may not exist or may be a CD or ZIP
    if os.system('[ -e /proc/ide/'+s+'/media ]') == 0:  # YESssss, we have a drive
        input1=open('/proc/ide/'+s+'/media','r')
        hdmedia=input1.readline()

        if string.find(hdmedia, "disk") == 0 :       # And by golly, it is a hard disk drive
            speedax=speedah[:]                       # Just reset the speed array for the next disk
            os.system('hdparm -i /dev/'+s+' > /etc/idedrives/.drives/parm'+s )    
            input2 = open('/etc/idedrives/.drives/parm'+s,'r')
            l=input2.readlines()
            l=l[3:]
            driveid0=string.split(l[0],',')
            driveid1=string.split(driveid0[0],'=')
            driveid=driveid1[1]
            lower=''
            upper=''
            
            tolerance=[]
            tct=0
            for g in l:
                
                if string.find(g,'MaxMultSec') > -1:
                    d1=string.split(g,',')
                    d2=string.split(d1[2],'=')
                    MaxMultSec=int(d2[1])
                    d3=string.split(d1[1],'=')
                    e1=d3[1][:-2]
                    BuffSize=int(e1)
                    continue
                if string.find(g,'PIO modes') > -1:
                    d1=string.split(g, ' ')
                    npio=len(d1)-2
                    continue
                if string.find(g,'DMA modes') > -1:
                    d1=string.split(g, ' ')
                    ndma=len(d1)-2
                    continue
                #         At this point, driva analysis needs either an accomplishes line OR
                #         a database lookup--we are using DB lookups, first preference and
                #         accomplishes lines second--note this requires a recent hdparm prog
                #         as in something that talks to kernel 2.4
                if string.find(g,'Drive Supp') > -1:
                    dx=string.split(g,':')
                    tolerance=string.split(dx[2],' ') # what the drive claims to support

                    tolerance=tolerance[1:-1]
                    tct=len(tolerance)
                    continue
            #
            #  OK all done with what we need from hdparm -i
            #
            #  Let's match with the DB if possible
            #
            # hack for damnable WDs that keep changing their ID format
            if string.find(driveid,'WDC ') > -1:
                driveid==driveid[4:]
            
            #
            #        End of hack to slice off the extraneous WDC_ that some WD drives report
            driveid=string.replace(string.upper(driveid),' ','')
            driveid=string.replace(driveid,'-','')
            #
            count=0
            for drive in drivedb:
                xdrive=string.split(drive)
                                                                   
                result=string.find(beta[count],driveid)
                if result>-1 :
                    lower=xdrive[1]
                    upper=xdrive[2]
                    print
                    print 'Drive '+driveid+' found in database '
                    print 'upper is '+upper
                    print 'lower is '+lower
                    print
                    break
                count=count+1
            if result <0 :
                print
                print 'Drive '+driveid+' not found in database'
                print 'Attempting to use capabilities line in hdparm -i'
                print 'Please report driveid to hdd@linuxtests.org'
                print
                if len(tolerance) == 0:
                        print 'Capabilities line not available this drive model'
                        print 'No optimizations possible this program'
                        print
                        continue
                if string.find(driveid,'WD') == -1:  #Damn those WD drives we don't trust them.
                    
                    if string.find(tolerance[0][:1],' ') >= 0:
                        tolerance[0]=tolerance[0][1:]
                    if string.find(tolerance[0],'ATA-1') >= 0:
                        lower='DMA'
                    elif string.find(tolerance[0],'ATA-2') >= 0:
                        lower='UDMA'
                    elif string.find(tolerance[0],'ATA-3') >= 0:
                        lower='UDMA2'
                    else:
                        print 'Capabilities line unusable--defaulting to low setting'
                        print
                        lower='PIO2'
                    if tct>1 :   
                        yy=tolerance[:-1][0]
                    else:
                        yy=lower
                    if string.find(yy,'ATA-6') >= 0 or string.find(yy,'ATA-5') >= 0:
                        upper='UDMA5'
                    elif string.find(yy,'ATA-4')>=0:
                        upper='UDMA4'
                    elif string.find(yy,'ATA-3')>=0  or string.find(yy,'ATA-2')>=0:
                        upper='UDMA2'
                    elif string.find(yy,'ATA-1')>=0:
                        upper='UDMA'
                    else:
                        print 'Capabilities line unusable--defaulting to low setting'
                        print
                        upper='PIO4'
                    print 'Settings are:'
                    print lower
                    print 'to'
                    print upper
                else:  #the accursed drive
                    print 'Western Digital Unknown Drive Detected'
                    print 'defaulting to UDMA2'
                    print
                    lower='UDMA2'
                    upper='UDMA2'
                    
            #  Now we have upper and lower limits
            if string.find(lower[:3],'PIO') > -1:
                print 'Optimization will not provide significant improvement for '+s
                print 'Proceeding to next drive'
                print
                continue
            else:
                mxct=0
                axct=0
                if BuffSize == 0:
                    maxmxct=0
                    maxaxct=0
                elif BuffSize > 0 and BuffSize < 129:
                    maxmxct=2
                    maxaxct=2
                elif BuffSize >128 and BuffSize <257:
                    maxmxct=4
                    maxaxct=4
                elif BuffSize < 1024 and BuffSize > 256:
                    maxmxct=8
                    maxaxct=8
                else:
                    maxmxct=16
                    maxaxct=8
                if MaxMultSec < maxmxct:
                    maxmxct=MaxMultSec
                SS=dmadic.keys()
                S=SS.sort()
                T=[]
                xcessp=''
                speedbase=speedtest(s)
                if seeker != 0:
                    print 'Default settings are too high for subject drive--errors encountered'
                    print
                    
                    speedbase=0
                print tolerance

                print SS
                print lower, upper
                
                for qq in SS:
                    if qq >= lower and qq <= upper:
                        T=T+[qq]
                print T
                
                print 'Continue(Y/n)?'
                answer=raw_input()
                if len(answer)==0:
                    answer='y'
                if answer[:1] == 'Y':
                    answer='y'
                if answer[:1] != 'y':
                    sys.exit(0)
                snogo=0
                for qq in T:
                    protocol=dmadic[qq]
                    speedcurrent=0
                    for readah in range(4):
                        for mult in range(4):
                            tstread=lookah[readah]
                            tstmult=multah[mult]
                            if tstread > maxaxct or tstmult > maxmxct:
                                xindex=dmarrdic[protocol]
                                scorespeed(xindex,readah,mult,0.0)
                                continue
                            xindex=dmarrdic[protocol]
                            xcess=parmset(s,1,1,1,tstmult,tstread,protocol)
                            print
                            print 'Testing '+s+' with readahead '+repr(tstread)+' multcount '+repr(tstmult)+' and protocol '+repr(dmardic[xindex])
                            print 'Default settings had speed '+repr(speedbase)
                            print
                            speedcurrent=speedtest(s)
                            if seeker!=0:
                                print
                                print 'Errors encountered--test terminated this drive'
                                print
                                if xcessp != '':
                                    os.system(xcessp)    # resets to last good setting so we can continue with next drive
                                erhandler(s,scompanion)
                                snogo=1
                                break
                            scorespeed(xindex,readah,mult,speedcurrent)
                            print 'Achieved '+repr(speedcurrent)
                            continue
                        if snogo==1:
                            break
                    if snogo==1:
                        break
                    
                    xb=bespeed()
                    Lset=xb[0]
                    Mset=xb[1]
                    Xset=xb[2]
                    protocolx=dmardic[Xset]
                    protocol=dmadic[protocolx]
                    print 'Best speed at '+repr(Mset)+' mult and lookahead '+repr(Lset)+' was '+repr(xb[3])
                    print 'for DMA setting '+repr(protocolx)    
                    xcess=parmset(s,1,1,1,Mset,Lset,protocol)
                    xcessp=xcess[:]
                    salt(xcess)
                    speedbase=speedtest(s)
                xb=bespeed()
                Lset=xb[0]
                Mset=xb[1]
                Xset=xb[2]
                protocolx=dmardic[Xset]
                protocol=dmadic[protocolx]
                xcess=parmset(s,1,1,1,Mset,Lset,protocol)
                salt(xcess)
                speedlast=speedtest(s)
                print 'Best speed at '+repr(Mset)+' mult and lookahead '+repr(Lset)+' was '+repr(speedlast)
                print 'for DMA setting '+repr(protocolx)    
                print
                print 'For drive '+s+', we had a default of'
                print repr(speedbase)+ 'Mb/sec, and we finished'
                print 'with a settting of '+repr(speedlast) +' Mb/sec'
                print
        # Go to the next drive
    #If we have one
# keep goint till we are done with all possible IDEs
os.system('echo "echo" >> /etc/rc.d/init.d/drivetune')
