#include "userlist.h"
#include "util.h"


#include <stdlib.h>
#include <utmp.h>
#include <errno.h>
#include <string.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <time.h>
#include <ctype.h>
#include <signal.h>

#define TRUE 1
#define FALSE 0
#define MIN(a,b) ((a) < (b) ? (a) : (b))

static void getSessionData( struct user * user );

static void fixString( char *buffer, int maxLen )
{
    while ( *buffer != '\0' && maxLen-- > 1 ) {
        if ( *buffer == ' ' || *buffer == '\t' || *buffer == '\r' || *buffer == '\n' ) {
            *buffer = '_';
        }
        buffer++;
    }
    *buffer = '\0';
}

static time_t lastActivity( char *line ) {
    struct stat s;
    char dev[40];
    snprintf( dev, sizeof(dev), "/dev/%s", line );
    if ( stat( dev, &s ) != 0 || s.st_atime == 0 )
        return -1;
    return s.st_atime;
}

int getUserList( struct user *outbuf, int size )
{
    int fh = open( "/run/utmp", O_RDONLY );
    if ( fh == -1 )
        return 0;
    // Mark
    int deadzone = 0;
    for ( int i = 0; i < size; ++i ) {
        if ( outbuf[i].user[0] != '\0' ) {
            deadzone = i + 1;
        }
        outbuf[i].lastOnline = outbuf[i].online;
        outbuf[i].online = false;
    }
    struct utmp buffer[100];
    ssize_t len = read( fh, buffer, sizeof(buffer) );
    int num = len / sizeof(struct utmp);
    for ( int i = 0; i < num; ++i ) {
        struct utmp *u = buffer + i;
        if ( u->ut_type != USER_PROCESS )
            continue; // Not user session
        if ( u->ut_user[0] == '\0' )
            continue; // Empty username!?
        // Check if alive
        if ( kill( u->ut_pid, 0 ) != 0 )
            continue;
        // Make sure these contain anything
        if ( u->ut_line[0] == '\0' ) {
            snprintf( u->ut_line, UT_LINESIZE, "-" );
        }
        if ( u->ut_host[0] == '\0' ) {
            snprintf( u->ut_host, UT_HOSTSIZE, "-" );
        }
        // No spaces, null terminated
        fixString( u->ut_user, UT_NAMESIZE );
        fixString( u->ut_line, UT_LINESIZE );
        fixString( u->ut_host, UT_HOSTSIZE );
        int use = -1;
        // Skip pts sessions within x (xterm)
        if ( u->ut_host[0] == ':' && isdigit( u->ut_host[1] ) ) {
            for ( int j = 0; j < deadzone; ++j ) {
                if ( outbuf[j].user[0] != '\0' && strcmp( outbuf[j].display, u->ut_host ) == 0 ) {
                    if ( outbuf[j].sessionLeader == u->ut_pid
                            && strncmp( outbuf[j].user, u->ut_user, UT_NAMESIZE ) == 0
                            && strncmp( outbuf[j].device, u->ut_line, UT_LINESIZE ) == 0 ) {
                        use = j;
                        break;
                    }
                    bool oldPts = strncmp( outbuf[j].device, "pts", 3 ) == 0;
                    bool newTty = strncmp( u->ut_line, "tty", 3 ) == 0;
                    if ( oldPts && newTty ) {
                        use = j; // Replace pts entry with tty entry
                        outbuf[j].user[0] = '\0';
                    } else if ( !newTty ) {
                        use = -2; // Ignore non-tty entry if we have another one for that display
                    }
                    break;
                }
            }
            if ( use == -2 )
                continue; // Ignore this entry, already taken care of
        }
        if ( use == -1 ) {
            // OK so far, find slot in outbuf (matching one, first free one if not found)
            for ( int j = 0; j < deadzone; ++j ) {
                if ( outbuf[j].user[0] == '\0' ) {
                    if ( use == -1 ) {
                        use = j;
                    }
                } else if ( outbuf[j].sessionLeader == u->ut_pid
                        && strncmp( outbuf[j].user, u->ut_user, UT_NAMESIZE ) == 0
                        && strncmp( outbuf[j].device, u->ut_line, UT_LINESIZE ) == 0 ) {
                    use = j;
                    break;
                }
            }
        }
        if ( use == -1 ) { // Not found and no free slot, expand if possible
            if ( deadzone >= size ) // No more slots :-/
                continue;
            use = deadzone++;
        }
        if ( outbuf[use].user[0] == '\0' ) {
            // New entry, init
            memset( &outbuf[use], 0, sizeof(outbuf[use]) );
            snprintf( outbuf[use].user, STRLEN, "%.*s",
                    (int)MIN( UT_NAMESIZE, STRLEN-1 ), u->ut_user );
            snprintf( outbuf[use].device, STRLEN, "%.*s",
                    (int)MIN( UT_LINESIZE, STRLEN-1 ), u->ut_line );
            if ( u->ut_host[0] == ':' && isdigit( u->ut_host[1] ) ) {
                snprintf( outbuf[use].display, STRLEN, "%.*s",
                        (int)MIN( UT_HOSTSIZE, STRLEN-1 ), u->ut_host );
            }
            outbuf[use].sessionLeader = u->ut_pid;
            getSessionData( &outbuf[use] );
            printf( "New Session: '%s' on '%s', Display '%s', logind session '%s', login pid '%d', utmp pid '%d'.\n", outbuf[use].user, outbuf[use].device, outbuf[use].display, outbuf[use].sessionName, (int)outbuf[use].sessionHead, (int)outbuf[use].sessionLeader );
        } else if ( outbuf[use].sessionName[0] == '\0' && outbuf[use].loginctlFails < 5 ) {
            getSessionData( &outbuf[use] );
            outbuf[use].loginctlFails++;
        }
        outbuf[use].online = true;
        // Reset offset if timestamp changed
        // but ONLY if this isn't a known X session
        if ( outbuf[use].display[0] == '\0' ) {
            const time_t la = lastActivity( outbuf[use].device );
            if ( outbuf[use].lastActivity != la ) {
                outbuf[use].lastActivity = la;
                outbuf[use].lastActivityOffset = 0;
            }
        }
    }
    close( fh );
    // Compact
    for ( int i = 0; i < deadzone; ++i ) {
        if ( outbuf[i].online || outbuf[i].lastOnline )
            continue; // In use, or just finished
        do {
            deadzone--;
            if ( outbuf[deadzone].online || outbuf[deadzone].lastOnline ) {
                outbuf[i] = outbuf[deadzone];
                outbuf[deadzone].user[0] = '\0';
                break;
            }
        } while ( deadzone > i );
    }
    return deadzone;
}

static void getSessionData( struct user * user )
{
    char buffer[500];
    snprintf( buffer, sizeof(buffer), "/proc/%d/cgroup", (int)user->sessionLeader );
    FILE *fh = fopen( buffer, "r" );
    if ( fh == NULL ) {
        fprintf( stderr, "getSessionData: Cannot open %s to get logind session name\n", buffer );
        return;
    }
    // Read from cgroup file
    user->sessionName[0] = '\0';
    char sessionName[40] = "";
    while ( fgets( buffer, sizeof(buffer), fh ) != NULL ) {
        const char *p = strstr( buffer, ":name=systemd:" );
        if ( p == NULL ) {
            p = buffer;
        }
        p = strstr( p + 14, "/session-" );
        if ( p == NULL )
            continue;
        copy( p + 9, sessionName, sizeof(sessionName), ".\n" );
        break;
    }
    fclose( fh );
    if ( sessionName[0] == '\0' ) {
        fprintf( stderr, "getSessionData: Cannot get logind session id from %d cgroup file.\n", (int)user->sessionLeader );
        return;
    }
    // Extract session leader pid from loginctl - it's often different from what utmp says
    int fds[2];
    if ( pipe( fds ) == -1 ) {
        perror( "getSessionData: pipe for loginctl failed" );
        return;
    }
    if ( doublefork() ) {
        close( fds[0] );
        redirect( -1, fds[1] );
        run( false, "loginctl", "-p", "Name", "-p", "Class", "-p", "Leader", "show-session", sessionName, (char*)NULL );
    }
    close( fds[1] );
    char lcName[100] = "", lcClass[100] = "", lcLeader[100] = "";
    if ( ! waitRead( fds[0], 2000 ) ) {
        fprintf( stderr, "getSessionData: loginctl timed out.\n" );
    } else {
        char output[4000] = "\n";
        char *pos = output + 1, *end = output + sizeof(output) - 1;
        do {
            ssize_t ret = read( fds[0], pos, end - pos );
            if ( ret == -1 ) {
                if ( errno == EINTR )
                    continue;
                perror( "getSessionData: read() from loginctl failed" );
                break;
            }
            if ( ret == 0 )
                break;
            pos += ret;
        } while ( waitRead( fds[0], 50 ) && pos < end );
        *pos = '\0';
        // Find strings
        pos = strstr( output, "\nName=" );
        if ( pos != NULL ) {
            copy( pos + 6, lcName, sizeof(lcName), "\n \t" );
        }
        pos = strstr( output, "\nClass=" );
        if ( pos != NULL ) {
            copy( pos + 7, lcClass, sizeof(lcClass), "\n \t" );
        }
        pos = strstr( output, "\nLeader=" );
        if ( pos != NULL ) {
            copy( pos + 8, lcLeader, sizeof(lcLeader), "\n \t" );
        }
    }
    close( fds[0] );
    // Eval what we got - only honor output if it makes sense
    if ( *lcName == '\0' ) {
        fprintf( stderr, "getSessionData: Could not get user name for session %s\n", sessionName );
        return;
    }
    if ( strcmp( lcName, user->user ) != 0 ) {
        fprintf( stderr, "getSessionData: Sanity check failed. Session %s belongs to %s, but expected %s.\n", sessionName, lcName, user->user );
        return;
    }
    if ( *lcClass != '\0' && strcmp( lcClass, "user" ) != 0 ) {
        fprintf( stderr, "killSession: Sanity check failed. Session %s has Class=%s, but expected user.\n", sessionName, lcClass );
        return;
    }
    // All good
    snprintf( user->sessionName, sizeof(user->sessionName), "%s", sessionName );
    user->sessionHead = atol( lcLeader );
    if ( user->sessionHead < 10 ) {
        fprintf( stderr, "getSessionData: Nonsensical Leader=%s reported by loginctl for session %s.\n", lcLeader, sessionName );
        user->sessionHead = 0;
    }
}