#!/usr/bin/perl

# needrestart-session - check for processes need to be restarted in user sessions
#
# Authors:
#   Thomas Liske <thomas@fiasko-nw.net>
#
# Copyright Holder:
#   2014 - 2015 (C) Thomas Liske [http://fiasko-nw.net/~thomas/]
#
# License:
#   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.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this package; if not, write to the Free Software
#   Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301 USA
#

use warnings;
use strict;

use constant {
    NRX11_ICO_FILE => q(/usr/share/needrestart-session/needrestart.xpm),
};

use NeedRestart::Utils;
use Wx qw(wxDefaultPosition wxDefaultSize wxPD_ELAPSED_TIME wxLC_REPORT wxALL wxTOP wxBOTTOM wxLIST_AUTOSIZE wxHORIZONTAL wxALIGN_CENTER wxEXPAND wxID_CLOSE wxID_REFRESH wxUSER_ATTENTION_INFO wxST_ELLIPSIZE_END wxBITMAP_TYPE_XPM);
use Wx::Event qw(EVT_BUTTON EVT_LIST_ITEM_ACTIVATED EVT_CLOSE EVT_TIMER);
use List::Util qw(max);
use Proc::ProcessTable;
use Cwd qw(realpath);

# Re-exec on SIGUSR1 - used by needrestart-dbus-session to trigger a
# refresh of needrestart-session
$SIG{USR1} = sub {
    exec($0, @ARGV);
};

# Don't show any progressbar nor 'Nothing found...' message if there
# was a parameter supplied (as needrestart-dbus-session does).
my $runsilent = scalar @ARGV;

# get current session
my $csession;
if(open(HCGROUP, qq(/proc/$$/cgroup))) {
    map {
	chomp;
	my ($id, $type, $value) = split(/:/);
	$csession = $2 if($type eq q(name=systemd) && $value =~ m@/user-(\d+).slice/(session-\d+).scope@);
    } <HCGROUP>;
    close(HCGROUP);
}
my @plist;

# probe needrestart's PolicyKit integration
my $haspk = `pkaction --action-id net.fiasko-nw.needrestart 2> /dev/null`;

# prepare GUI stuff
my $app = Wx::SimpleApp->new;
$app->SetAppName('needrestart-session');

# while running in Wx's mainloop perl won't get any signals, therefore
# we just run a timer event every 5 seconds to make signal handling
# deterministic
my $timer = Wx::Timer->new( $app );
$timer->Start( 5000 );
EVT_TIMER($app, -1, sub { 1; });

my $frame = Wx::Frame->new(undef, -1, 'needrestart-session');
$frame->SetIcon(Wx::Icon->new(NRX11_ICO_FILE, wxBITMAP_TYPE_XPM))
    if(-r NRX11_ICO_FILE);

my $sApps = Wx::StaticText->new($frame, -1, 'Some of your processes need to be restarted. You might need to re-login.', wxDefaultPosition, wxDefaultSize, wxST_ELLIPSIZE_END);
my $lApps = Wx::ListCtrl->new($frame, -1, wxDefaultPosition, [640, 480], wxLC_REPORT);
$lApps->InsertColumn(0, "Application");
$lApps->InsertColumn(1, "Session Type");
$lApps->InsertColumn(2, "Session Id");
$lApps->InsertColumn(3, "PIDs");

&RefreshAppsList;

$lApps->SetColumnWidth(0, wxLIST_AUTOSIZE);
$lApps->SetColumnWidth(1, wxLIST_AUTOSIZE);
$lApps->SetColumnWidth(2, wxLIST_AUTOSIZE);
$lApps->SetColumnWidth(3, wxLIST_AUTOSIZE);

my $bsButtons = Wx::BoxSizer->new(wxHORIZONTAL);
my $bRoot = Wx::Button->new($frame, -1, '&System Services');
$bRoot->Enable($haspk);
my $bRecheck = Wx::Button->new($frame, wxID_REFRESH);
my $bClose = Wx::Button->new($frame, wxID_CLOSE);

my $sDClick = Wx::StaticText->new($frame, -1, '*) Double click to bring application window to front.');

my $fgs = Wx::FlexGridSizer->new(6, 1, 0, 0);
$fgs->Add($sApps, 0, wxEXPAND | wxTOP | wxBOTTOM, 4);
$fgs->Add($lApps, 0, wxEXPAND);
$fgs->Add($sDClick, 0, wxEXPAND | wxTOP | wxBOTTOM, 4);
$bsButtons->Add($bRoot, 0, wxALL, 4);
$bsButtons->Add($bRecheck, 0, wxALL, 4);
$bsButtons->Add($bClose, 0, wxALL, 4);
$fgs->Add($bsButtons, 0, wxALIGN_CENTER);

$fgs->AddGrowableCol(0);
$fgs->AddGrowableRow(1);

$frame->SetSizerAndFit($fgs);
$frame->SetAutoLayout(1);
$frame->Centre();

# bind event handlers
EVT_BUTTON($frame, $bRoot, \&OnRunAsRoot);
EVT_BUTTON($frame, $bRecheck, \&OnRecheck);
EVT_BUTTON($frame, $bClose, \&OnClose);
EVT_CLOSE($frame, \&OnClose);
EVT_LIST_ITEM_ACTIVATED($frame, $lApps, \&OnAppsDClick);

# show main window and enter event loop
$frame->Show;
$app->MainLoop;

sub RefreshAppsList {
    my $frProgress = ($runsilent ? undef : Wx::ProgressDialog->new("", "", 4, undef, wxPD_ELAPSED_TIME));

    # get user processes required to be restarted
    my %fnames;
    my $fh = nr_fork_pipe(0, qw(needrestart -b));
    $frProgress->SetTitle(q(Scanning processes...)) if(defined($frProgress));
    while(<$fh>) {
	chomp;

	next unless(/^NEEDRESTART-PID: (.+)=([\d,]+)$/);

	@{ $fnames{$1} } = split(/,/, $2);
	if(defined($frProgress)) {
	    $frProgress->Update(1, $1);
	    $frProgress->Pulse;
	}
    }
    close($fh);

    # terminate if no orphan process is found
    if(scalar keys %fnames == 0) {
	$frProgress->Destroy if(defined($frProgress));

	unless($runsilent) {
	    my $msg = Wx::MessageDialog->new($frame, "None of your processes need to be restarted.", "Nothing found...");
	    $msg->ShowModal;
	    $msg->Destroy;
	}
	exit;
    }
    
    # get list of pid => window mapping
    $frProgress->SetTitle(q(Scanning X11 windows...)) if(defined($frProgress));
    my %ptable = map {$_->pid => $_} @{ new Proc::ProcessTable(enable_ttys => 1)->table };
    my %windows;
    my %wtitles;
    $fh = nr_fork_pipe(0, qw(wmctrl -l -p));
    while(<$fh>) {
	chomp;

	next unless(/^(0x[\da-f]+) +-?\d+ +(\d+) +(.+)/);

	$windows{$2} = $1;
	$wtitles{$2} = $3;
	if(defined($frProgress)) {
	    $frProgress->Update(2, $ptable{$2}->{fname});
	    $frProgress->Pulse;
	}
    }
    close($fh);

    # build process list from `needrestart` and `wmctrl` outputs
    $frProgress->SetTitle(q(Updating list...)) if(defined($frProgress));
    my %plist;
    foreach my $fname (keys %fnames) {
	foreach my $pid (@{$fnames{$fname}}) {
	    my $session = '?';
	    my $type = '';
	    if(exists($ptable{$pid})) {
		if($ptable{$pid}->{ttydev} ne '') {
		    $session = realpath( $ptable{$pid}->{ttydev} );
		}
		elsif(open(HCGROUP, qq(/proc/$pid/cgroup))) {
		    map {
			chomp;
			my ($id, $type, $value) = split(/:/);
			$session = $2 if($type eq q(name=systemd) && $value =~ m@/user-(\d+).slice/(session-\d+).scope@);
		    } <HCGROUP>;
		    close(HCGROUP);
		}

		if(defined($csession) && $session ne '?') {
		    $type = ($session eq $csession ? 'current' : 'foreign');
		}
	    }

	    if(grep { $_ == $pid } keys %windows) {
		push(@{ $plist{$fname}->{plist} }, {
		    fname => $fname,
		    PIDs => [$pid],
		    sessid => $session,
		    sessty => $type,
		    winid => $windows{$pid},
		    winti => $wtitles{$pid},
		});
	    }
	    elsif(exists($plist{$fname}->{sessions}->{$session})) {
		push(@{ $plist{$fname}->{sessions}->{$session}->{PIDs} }, $pid);
	    }
	    else {
		$plist{$fname}->{sessions}->{$session} = {
		    fname => $fname,
		    PIDs => [$pid],
		    sessid => $session,
		    sessty => $type,
		};
	    }
	}

	if(defined($frProgress)) {
	    $frProgress->Update(3, $fname);
	    $frProgress->Pulse;
	}
    }

    @plist = (
	(map { (exists($plist{$_}->{plist}) ? @{ $plist{$_}->{plist} } : ()); } keys %plist),
	(map {
	    # sort PIDs ascending
	    foreach my $sess (keys %{ $plist{$_}->{sessions} }) {
		$plist{$_}->{sessions}->{$sess}->{PIDs} = [sort {$a <=> $b} @{ $plist{$_}->{sessions}->{$sess}->{PIDs} }];
	    }
	    (values %{ $plist{$_}->{sessions} });
	 } keys %plist),
	);

    # sort processes by session, fname and pid
    @plist = sort {
	$a->{sessty} cmp $b->{sessty} ||
	    $a->{sessid} cmp $b->{sessid} ||
	    $a->{fname} cmp $b->{fname} ||
	    $a->{PIDs}->[0] <=> $b->{PIDs}->[0]
    } @plist;

    # (re)fill app list control
    $lApps->DeleteAllItems;
    my $pos = 0;
    foreach my $data (@plist) {
	$lApps->InsertStringItem($pos, $data->{fname} . (exists($data->{winti}) ? '*' : ''));
	$lApps->SetItem($pos, 1, $data->{sessty});
	$lApps->SetItem($pos, 2, $data->{sessid});
	$lApps->SetItem($pos++, 3, join(', ', @{ $data->{PIDs} }));
    }

    $frProgress->Destroy if(defined($frProgress));

    # Don't keep silence if user hits 'Refresh'.
    $runsilent++;
}

# make selected window active
sub OnAppsDClick {
    my ($frame, $event) = @_;

    my $pid = fork();

    system(qw(wmctrl -i -a), $plist[$event->GetIndex]->{winid}) if(exists($plist[$event->GetIndex]->{winid}));
}

sub OnRunAsRoot {
    my %frontends = (
	kde => [qw(QtCore4 QtGui4)],
    );
    my $desktop = lc($ENV{XDG_CURRENT_DESKTOP});
    my $frontend = q(gnome);

    # use desktop specific frontend if available
    if(defined($desktop) && exists($frontends{$desktop})) {
	my $e = map { eval "require $_;"; if($@) {(1);} else {()}; } @{$frontends{$desktop}};
	$frontend = $desktop unless($e);
    }

    # run needrestart as root using PolicyKit and overwritten debconf frontend
    my $pid = fork();
    unless($pid) {
	exec(qw(pkexec /usr/sbin/needrestart -f), $frontend);
	die;
    }
}

sub OnRecheck {
    RefreshAppsList;
}

sub OnClose {
    $app->ExitMainLoop;
}
