#!/usr/local/bin/perl -w
#
# script to equally distribute mailboxes of a cyrus imapd on differnt 
# imap partitions
#
# by Dietmar Rieder <dietmar.rieder(at)tugraz.at>
#    Gernot Stocker <gernot.stocker(at)tugraz.at>  
#
#    (c) 2006
#
use strict;
use Cyrus::IMAP::Admin;
use Filesys::DiskFree;
use List::Util 'shuffle';

#
# configuration: please adapt it to your needs
#
my $imap_host       = 'localhost';
my $imap_port       = '1430';
my $imap_admin      = 'imap';
my $imap_config     = '/etc/imapd.conf';
my $auth_mech       = 'login';
my $imap_admin_pass = 'secretpassword';

# safety belt
#exit;

# define some variables
my %config;
my @partitions;


# generate any output immediately
$| = 1;

# read imapd config
read_config($imap_config);

# determine the defaultpartition name and filesystem
my $default_part_name = 'partition-' . $config{'defaultpartition'};
my $default_patrition_fs = $config{$default_part_name};

# get and set all imap partitions
get_partitions();


# connect to the cyrus imap server as imap admin
my $cyradm = Cyrus::IMAP::Admin->new($imap_host, $imap_port) or die "Can't connect to $imap_host.\n";
$cyradm->authenticate(-user => "$imap_admin", -password => $imap_admin_pass, -mechanism => $auth_mech) or die "Can't authenticate as $imap_admin.\n";

# get all user mailboxes (toplevel)
my @mailboxes = cyrus_get_mailboxes($cyradm, 'user.%');

# shuffle the mailbox list and move them to the partitions which have the most free space
foreach my $mb (shuffle(@mailboxes)) {
    my %mailbox_attribs = cyrus_get_attrib($cyradm, $mb);
    my $old_mailbox_part = $mailbox_attribs{'partition'};
    my $mailbox_tot_size = $mailbox_attribs{'size'};

    # calculate mailbox total size
    my $sub_filter = $mb . '.*';
    foreach my $sub_mb (cyrus_get_mailboxes($cyradm, $sub_filter)) {
	my %sub_mb_attribs = cyrus_get_attrib($cyradm, $sub_mb);
	$mailbox_tot_size += $sub_mb_attribs{'size'};
    }

    # determine the best partition with hopefully more free space
    my $destination_part = get_new_part($old_mailbox_part, $mailbox_tot_size);

    # move the mailbox to the new partition if it is a different one
    if($destination_part ne $old_mailbox_part) {
	print "moving mailbox \"" . $mb . "\" from partition \"" . $old_mailbox_part . "\" to partition \"" . $destination_part . "\" (size: " . $mailbox_tot_size . ")\n";
	# TODO: remove print statement in production
	# print "cyradm->renamemailbox($mb, $mb, $destination_part)\n";
	$cyradm->renamemailbox($mb, $mb, $destination_part);
	# TODO: remove comment
	if($cyradm->error) {
	    print "\nError moving mailbox \"" . $mb . "\" from partition \"" . $old_mailbox_part . "\" to partition \"" . $destination_part . "\" : " . $cyradm->error . "\n";
	    print "\tkeeping mailbox \"" . $mb . "\" on partition \"" . $old_mailbox_part . "\" (size: " . $mailbox_tot_size . ")\n";
	    next;
	}		    
    } else {
	print "keeping mailbox \"" . $mb . "\" on partition \"" . $old_mailbox_part . "\" (size: " . $mailbox_tot_size . ")\n";
    }
}

exit 0;
# END

#####################
# subroutines
##


# get the mailbox attributes
sub cyrus_get_attrib {
    my $cyradm = shift;
    my $mailbox = shift;
    
    my %attribs;
    
    my %mbox_info = $cyradm->getinfo($mailbox);
    foreach my $attrib (sort keys %mbox_info) {
	$attrib =~ /([^\/]*)$/;
	my $attrname = $1;
	$attribs{$attrname} = $mbox_info{$attrib};
    }
    return %attribs;
}

# get a list of all user mailboxes (toplevel)
sub cyrus_get_mailboxes {
    my $cyradm = shift;
    my $filter = shift;

    my @mailbox_list;
    
    my @mboxes = $cyradm->list($filter);
    my $mbox_count = @mboxes;
    for(my $i = 0; $i < $mbox_count; $i++) {
	#print $mboxes[$i][0] . "\n";
	push(@mailbox_list, $mboxes[$i][0]);
    }
    return @mailbox_list;
}

# read the imapd.conf file
sub read_config {
    my $cfg = shift;
        
    open(CFG, "<$cfg") || die("$!: $cfg");
    while(<CFG>) {
	chomp;
	next if $_ =~ /^#|^$/;
	my($key, $val) = split(/:/, $_);
	$key =~ s/\s//g;
	$val =~ s/\s//g;
	$config{$key} = $val;
    }
}

# get all partitions form the imapd.conf file
sub get_partitions {
    foreach my $key (keys(%config)) {
	if($key =~ /^partition/) {
	    my($dummy, $partition) = split(/\-/, $key);
	    push(@partitions, $partition);
	}
    }
}

# determin the best partition for moving a mailbox
sub get_new_part {
    my $old_part = shift;
    my $mbox_size = shift;
    
    # get the available free disk space
    my $handle = new Filesys::DiskFree;
    $handle->df();
    
    my $max_free_part = $old_part;
    my $old_part_name = 'partition-' . $old_part;
    my $old_part_fs = $config{$old_part_name};
    my $max = $handle->avail($old_part_fs);
    my $old_part_size = $max;

    # find the partition with the most free space
    foreach my $part (@partitions) {
	my $new_part_name = 'partition-' . $part;
	my $new_part_fs = $config{$new_part_name};
	
	if ($max < $handle->avail($new_part_fs)) {
	    $max = $handle->avail($new_part_fs);
	    $max_free_part = $part;
	}
    }
    undef($handle);

    # we just move to a new partition if the free space on it is
    # 2x as big as the mailbox itself
    if($max - $old_part_size < $mbox_size * 2) {
	return($old_part);
    }

    if($max_free_part ne $old_part) {
	return $max_free_part;
    }
    return($old_part);
}