####################################################################
# Windows 2000 Disk Quota Utility
my $version='0.87b'; my $verdate='09/28/00';
###################################################################
# use strict    #  <vlad ducks the tomatoes>
use Win32::OLE;

my @execcue=();			# List that that holds functions to be exec'd
my $verbose=0;
my $quickdel=0;			# Flag for quick deletion of quotas via a hack

my %quotasize=(			# quota limits in bytes
  "group1" => 209715200,	# 200MB
  "group2" => 314572800,	# 300MB
  "group3" => 576716800		# 500MB
);

my %quotathresh=(		# quota threshold in bytes
  "group1" => 209715200,	# 200MB
  "group2" => 314572800,	# 300MB
  "group3" => 576716800		# 500MB
);

my %argsnorm=(
  "qa" => ["QuotaAdd", "Add or Modify a user's Quota"],
  "qd" => ["QuotaDelete","Delete a Quota Entry"],
  "qv" => ["QuotaView","Display a User's Quota Information"],
  "qs" => ["QuotaStatus","Quota System Status Information"],
  "qc" => ["QuotaConfig","Quota System Configuration"]
);

my %argsdata=(
  "u" => ["username","User Account"],
  "g" => ["group","Quota Group for predefined size/thresholds"],
  "s" => ["size","Size of Quota Limit in MB"],
  "t" => ["thresh","Warning Threshold in MB"],
  "r" => ["domain","Windows Domain for User Account (Blank for LocalSystem)"],
  "m" => ["dqmode","Quota System Enforcement Mode (0=None,1=Soft,2=Hard)"],
  "l" => ["logfile","Name of Logfile if desired"],
  "d" => ["drive","Quota Drive (D or D:)"],
  "z" => ["dqsize","Default Quota Size in MB"],
  "e" => ["dqthresh","Default Threshold Size in MB"],
  "m" => ["loglimit","Limit Hit Logging (0 or 1)"],
  "c" => ["logthresh","Threshold Hit Logging (0 or 1)"],
  "o" => ["usrres","User Reolution Mode (0-No,1-Yes,2-Yes/Async)"],
);

####################################################################
# Argument Parsing
####################################################################
&Usage unless (@ARGV);
ARG: while (my $arg=lc(shift @ARGV)) {
  $arg=~s/^(-+|\/|\\)//;			 	# remove the -+ or / or \ 
  if ($arg eq 'w') {$^W=1;next ARG}			# turn on warnings
  if ($arg eq 'v') {$verbose=1;next ARG}		# turn on verbose
  if ($arg eq 'q') {$quickdel=1;next ARG}            # Quick delete hack (see &QuotaDelete)  
  if ($arg=~/^\?$|^h$|^help$/) {&Usage}
  foreach (keys %argsnorm) { 				# for each normal arg that doesn't have stuff after it..
    if ($arg eq $_) { 					
      push @execcue, $arg;				# push commands in order to the execution cue
      next ARG; 					# next arg
    }
  }
  foreach (sort keys %argsdata) { 				# for each special arg that has useful stuff after it..
    if ($arg eq $_) {						
      ${${$argsdata{$arg}}[0]}=shift @ARGV; 			# set the correponding varible to the next arg         
      next ARG;	 						
    }
  }
  print "\nERROR: $arg is not a valid option.\n";
  print "For help, type dq.pl -h\n\n";
  exit(0);
}
####################################################################
# Argument Prep & Sanity Checking
####################################################################
if (defined $logfile) {
  my $loghandle=select(LOG);$|=1;select($loghandle);
  open(LOG,">>$logfile");
}
if ($verbose) {
  my ($day,$mth,$year)=(localtime)[3,4,5];	# For logging or stdout prints
  $year+=1900;
  $mth++;
}
if ($drive=~/^(\w)$/) {
  $drive=$1.":";
}
elsif ((not defined $drive) or (not /^\w\:$/)) {
  print("ERROR: You need to specify a quota drive, like D:\n\n");
  &Usage  
}
else {
  print("ERROR: You need to specify a quota drive, like D:\n\n");
  &Usage  
}	
if (not defined $domain) {
  $domain=$ENV{'COMPUTERNAME'};	
  print "No domain specified with -d, using $domain.\n";
}
if (grep/qa|qd|qv/,@execcue) {
  if (not defined $username) {
    print "You must specify a username!\n";
    &Usage;
 }
}
if (grep/qc/,@execcue) {
  if (defined $dqsize) {$dqsize=(($dqsize*1024)*1024)}
  if (defined $dqthresh) {$dqthresh=(($dqthresh*1024)*1024)}
  if ((defined $loglimit) and (not $loglimit le 1)) {
    &LogPrint("Logging Options are 0 or 1 (Off or On");
    &Usage;
  }
  if ((defined $logthresh) and (not $logthresh le 1)) {
    &LogPrint("Logging Options are 0 or 1 (Off or On");
    &Usage;
  }
  if ((defined $usrres) and (not $usrres le 2)) {
    &LogPrint("User Resolution can be set to 0, 1 or 2");
    &Usage;
  }
  if ((defined $dqmode) and (not $dqmode le 2)) {
    &LogPrint("Enforcement Mode can be set to 0, 1 or 2");
    &Usage;
  }
}
if (grep/qa/,@execcue) {
  if (defined $group) {
    $size=$quotasize{$group};	
    $thresh=$quotathresh{$group};	
  }
  elsif ((defined $size) and (not defined $thresh)) {
    print "No threshold defined, using $size MB (specify with -t)\n";
    $size=(($size*1024)*1024);
    $thresh=$size;
  }
  elsif ((defined $size) and (defined $thresh)) {
    $size=(($size*1024)*1024);
    $thresh=(($thresh*1024)*1024);
  }
  else {
  LogPrint("You must specify the quota size (-s) in MB, or a group name with -g\n");
  &Usage;
  }
}
####################################################################
&Main;
####################################################################
sub Main {
  for (@execcue) {
    if(&{${$argsnorm{$_}}[0]}) {&LogPrint("$argsnorm{$_}[0] Success") if $verbose}  	# call the correponding sub
    else{&LogPrint("$argsnorm{$_}[0] Failed")}
  }
  close LOGFILE if (defined $logfile);
  exit(1);
}
###################################################################
# Quota - Init
###################################################################
sub QuotaInit {

  my $dq=Win32::OLE->CreateObject("Microsoft.DiskQuota.1") or &ErrPrint("Couldn't create new instance of Disk Quota Control: $!",1);
  &OLEErrChk("Creating User Quota Object",1);

  $dq->Initialize($drive,1);				 
  &OLEErrChk("Initializing User Quota Object");	
  
  return $dq		
}
###################################################################
# Quota - Status Info
###################################################################
sub QuotaStatus {
  my $dq=&QuotaInit;
 
  my $dqrebuild=$dq->QuotaFileRebuilding;			# For debugging to see if the Quota file is being
  &OLEErrChk("Checking Rebuild Status");			# rebuilt.  
  &LogPrint("Quota File Rebuilding? $dqrebuild");
											
  my $dqfilestate=$dq->QuotaFileIncomplete;			# Checks to see if the Quota File is incomplete.
  &OLEErrChk("Checking Quota State");
  &LogPrint("Quota File Incomplete? $dqfilestate");

  undef $dq;
  return 1;
}
###################################################################
# Quota - Config
###################################################################
sub QuotaConfig {
  my $dq=&QuotaInit;

  if (defined $dqmode) {
    $dq->{QuotaState}=$dqmode;
    &OLEErrChk("Setting Quotas to $dqmode for $drive");
  }		
  if (defined $dqsize) {
    $dq->{DefaultQuotaLimit}=$dqsize;
    &OLEErrChk("Setting Default Quota Size");
  }
  if (defined $dqthresh) {
    $dq->{DefaultQuotaThreshold}=$dqthresh;
    &OLEErrChk("Setting Default Quota Threshold");
  }
  if (defined $loglimit) {
    $dq->{LogQuotaLimit}=$loglimit;
    &OLEErrChk("Setting Limit Hit Logging");
  }
  if (defined $logthresh) {
    $dq->{LogQuotaThreshold}=$logthresh;
    &OLEErrChk("Setting Threshold Hit Logging");
  }
  if (defined $usrres) {					# While creating a quota entry, the Quota subsystem runs off
    $dq->{UserNameResolution}=$usrres;				# to resolve this SID. It seems damn slow.
    &OLEErrChk("Setting Resolution State to dqResolveNone");	# dbResolveAsync(2) wants to be a fancy way of not waiting, but
  }								# seems to be slow too. Even in small enviroments, this can be slow.
 
  my $dqdefsize=$dq->DefaultQuotaLimitText;
  &OLEErrChk("Checking Default Quota Size");
  &LogPrint("Default User Quota Size: $dqdefsize");
											
  my $dqdefthresh=$dq->DefaultQuotaThresholdText;
  &OLEErrChk("Checking Default Quota Threshold");
  &LogPrint("Default User Quota Threshold: $dqdefthresh");

  my $dqloglimit=$dq->LogQuotaLimit;
  &OLEErrChk("Checking Limit Hit Logging");
  &LogPrint("Limit Hit Logging: $dqloglimit");

  my $dqlogthresh=$dq->LogQuotaThreshold;
  &OLEErrChk("Checking Threshold Hit Logging");
  &LogPrint("Threshold Hit Logging: $dqlogthresh");

  my $dqstate=$dq->{QuotaState};
  &OLEErrChk("Checking Quota Enforce State");
  &LogPrint("Current Enforce State: $dqstate (0-Off,1-Soft,2-Hard)");
								
  my $dqres=$dq->{UserNameResolution};				
  &OLEErrChk("Setting Resolution State to dqResolveNone");	
  &LogPrint("User Resolution State: $dqres (0-No,1-Yes,2-Yes/Async)");			

  undef $dq;
  return 1;
}
###################################################################
# Quota - Create User Entry
###################################################################
sub QuotaAdd {
  my $hack=shift;
  my $dq=&QuotaInit;

  my $dqUser=$dq->AddUser("$domain\\$username");
  &OLEErrChk("Making Quota Entry for $username");

  $dqUser->{QuotaLimit}=$size;
  &OLEErrChk("Setting Quota Limit for $username ($size)");

  $dqUser->{QuotaThreshold}=$thresh;
  &OLEErrChk("Setting Threshold for $username at ($thresh)");

  undef $dq;
  return $dqUser if ($hack);
  undef $dqUser;
  return 1;
}
###################################################################
# Quota - Delete User Entry
###################################################################
sub QuotaDelete {
  my $dq=&QuotaInit;
  my $dqUser=&QuotaCreate(1) if $quickdel;			# The FindUser Method is horribly slow.
  my $dqUser=&QuotaFind if not $quickdel; 				# So we can hack it by doing a  create & delete

  $dq->DeleteUser($dqUser);					# dear god this sucks, we create & delete, if you want.
  &OLEErrChk("Deleting Quota Entry for $domain\\$username");	# It's a matter of 2 seconds vs. 10+ on big installs (500users+)
								# with the same end result.
  undef $dq;
  undef $dqUser;
  
  return 1;
}
###################################################################
# Quota - View User Entry
###################################################################
sub QuotaView {
  my $dq=&QuotaInit;
  my $dqUser=&QuotaFind;
  
  my $strsize="Quota Size for $username is " . (($dqUser->{QuotaLimit}/1024)/1024) . "MB";
  LogPrint($strsize);
  &OLEErrChk("Retieving Quota Limit for $username");

  my $strthresh="Quota Threshold for $username is " . (($dqUser->{QuotaThreshold}/1024)/1024) . "MB";
  LogPrint($strthresh);
  &OLEErrChk("Retieving Quota Threshold for $username");

  undef $dq;
  undef $dqUser;

  return 1;
}
###################################################################
# Quota - Find User Entry
###################################################################
sub QuotaFind {
  my $dq=&QuotaInit;

  my $dqUser=$dq->FindUser("$domain\\$username");
  my $errNum = Win32::OLE->LastError;
  if ($errNum != 0) {
    &LogPrint(Win32::FormatMessage($errNum));
    &LogPrint("User not found.");
    undef $dqUser;
    exit(0);
  }

  return $dqUser;
}
###################################################################
# OLE Error Check
###################################################################
sub OLEErrChk {
  my $oleaction=shift;
  my ($package, $filename, $line, $subname, $hasargs, $wantarray)=caller(1);	# Use caller to see where/why LogPrint was called from
  $subname=~s/main:://;
  my $errNum = Win32::OLE->LastError;
  if ($errNum==0) {&LogPrint("$oleaction was Successful") if ($verbose)}	# Log the successfule OLE action
  else {
    $errNum=~s/\n/ /g;						# rip out newlines
    die "$oleaction failed ($subname($line)): $errNum";
  }
  return 1;
}
###################################################################
# Log Print
###################################################################
sub LogPrint {
  my $message=shift;
  my $chatty=shift;
  
  if ($verbose) {
    my ($package, $filename, $line, $subname, $hasargs, $wantarray)=caller(1);	# Use caller to see where/why LogPrint was called from
    $subname=~s/main:://;
  	
    my ($sec,$min,$hour)=localtime;
    $chatty="$mth/$day/$year,$hour:$min:$sec,$subname($line),$domain,$username,";
  }

  print LOG "$chatty$message\n"  if (defined($logfile));
  print "$chatty$message\n";
}
###################################################################
# Usage
###################################################################
sub Usage {
  print "Version: $version   Date: $verdate\n";
  print "Usage: dq.pl -u [user] -s [size] [action]\n";
  foreach (sort keys %argsnorm) {
    print "-$_   /$_  \t$argsnorm{$_}[1]\n";	
  }
  print "\n";
  foreach (sort keys %argsdata) {
    print "-$_ /$_ $argsdata{$_}[1]\n";	
  }
  exit(0);
}