rmps/agent/agent.c

/*
 * agent.c
 *
 * Copyright (C) 2018 by Bogomil Vasilev <b.vasilev@smirky.net>
 *
 * This file is part of Remote Management and Provisioning System (RMPS).
 *
 * RMPS 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.
 *
 * RMPS 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 RMPS.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <fcntl.h>
#include <unistd.h>
#include <pthread.h>
#include <string.h>
#include <signal.h>
#include "job.h"
#include "../src/protocol.h"
#include "agent_ssl.h"

#define FAIL -1

static SSL_CTX *ctx;
static SSL *ssl;
static int server;
static struct job_args *args;
static pthread_t *job_thread;
static void agent_shutdown(void);
static void signal_handler(int sig);

static short get_job_slot();

static short get_job_slot()
{
	unsigned short i;
	for (i = 0; i < MAX_AGENT_JOBS; i++)
		if (args[i].slot != FULL)
			return i;
	return -1;
}

static short find_job(unsigned id)
{
	unsigned short i;
	for (i = 0; i < MAX_AGENT_JOBS; i++)
		if (args[i].buf.meta.id == id)
			return i;
	return -1;
}

static void signal_handler(int sig)
{
	switch (sig) {
	case SIGHUP:
		printf("Received SIGHUP signal. Ignoring...\n");
		break;
	case SIGINT:
		printf("Received SIGINT signal!\n");
		agent_shutdown();
		printf("Agent has been stopped properly.\n");
		_exit(EXIT_SUCCESS);
	case SIGTERM:
		printf("Received SIGTERM signal!\n");
		agent_shutdown();
		printf("Agent has been stopped properly.\n");
		_exit(EXIT_SUCCESS);
	default:
		printf("Unhandled signal %s", strsignal(sig));
		break;
	}
}

static void agent_shutdown(void)
{
	printf("Shutting down agent...\n");
	if (ssl && !SSL_get_shutdown(ssl))
		SSL_shutdown(ssl);
	if (ssl) SSL_free(ssl);
	if (ctx) SSL_CTX_free(ctx);	/* release context */
	close(server);			/* close socket */
	_exit(EXIT_SUCCESS);
}

static void set_env(void)
{
	struct sigaction new_sigaction;
	sigset_t new_sigset;

	/* Set signal mask - signals we want to block */
	sigemptyset(&new_sigset);
	sigaddset(&new_sigset, SIGCHLD);  /* ignore child */
	sigaddset(&new_sigset, SIGTSTP);  /* ignore Tty stop signals */
	sigaddset(&new_sigset, SIGTTOU);  /* ignore Tty background writes */
	sigaddset(&new_sigset, SIGTTIN);  /* ignore Tty background reads */
	sigprocmask(SIG_BLOCK, &new_sigset, NULL); /* Block above signals */

	/* Set up a signal handler */
	new_sigaction.sa_handler = signal_handler;
	sigemptyset(&new_sigaction.sa_mask);
	new_sigaction.sa_flags = 0;

	/* Signals to handle */
	sigaction(SIGHUP, &new_sigaction, NULL);  /* catch hangup signal */
	sigaction(SIGTERM, &new_sigaction, NULL); /* catch term signal */
	sigaction(SIGINT, &new_sigaction, NULL);  /* catch interrupt signal */
	signal(SIGPIPE, SIG_IGN); /* prevent crashing from bad writes */
}

int main(int count, char *strings[])
{
	int bytes, i, err;
	char *hostname, *portnum;

	if (count != 6) {
		printf("usage: %s <hostname> <portnum> <key> <cert> <CA>\n", strings[0]);
		_exit(EXIT_SUCCESS);
	}
	hostname=strings[1];
	portnum=strings[2];
	if ((ctx = init_ctx()) == NULL)
		_exit(EXIT_FAILURE);

	while (1) {
		server = connect_to_rmps(hostname, atoi(portnum));
		if (!server) {
			fprintf(stderr, "Failed to connect to RMPS on %s:%d - %s\n", hostname, atoi(portnum), strerror(errno));
			printf("Retrying...\n");
			sleep(5);
			continue;
		}
		load_certs(ctx, strings[3], strings[4], strings[5]);
		ssl = SSL_new(ctx);		/* create new SSL connection state */
		SSL_set_fd(ssl, server);	/* attach the socket descriptor */

		if (SSL_connect(ssl) == FAIL) {	/* perform the connection */
			ERR_print_errors_fp(stderr);
			close(server);
			SSL_CTX_free(ctx);
			_exit(EXIT_FAILURE);
		}
		printf("Connected with %s encryption\n", SSL_get_cipher(ssl));
		show_certs(ssl);
		set_env();
		atexit(agent_shutdown);
		if (!(args = calloc(1, sizeof(*args) * MAX_AGENT_JOBS))) {
			fprintf( stderr,
				 "Failed to calloc() %d bytes for job_args! Exiting...\n",
				 (int)sizeof(struct job_args) * MAX_AGENT_JOBS );
			SSL_shutdown(ssl);
			SSL_free(ssl);
			close(server);
			SSL_CTX_free(ctx);
			_exit(EXIT_FAILURE);
		}
		if (!(job_thread = calloc(1, sizeof(*job_thread) * MAX_AGENT_JOBS))) {
			fprintf( stderr,
				 "Failed to calloc() %d bytes for job_threads! Exiting...\n",
				 (int)sizeof(pthread_t) * MAX_AGENT_JOBS );
			SSL_shutdown(ssl);
			SSL_free(ssl);
			close(server);
			SSL_CTX_free(ctx);
			free(args);
			_exit(EXIT_FAILURE);
		}
		for (i = 0; i < MAX_AGENT_JOBS; i++) {
			args[i].slot = FREE;
			args[i].ssl = ssl;
		}

		do {
			struct msg_t buf;
			memset(&buf, 0, sizeof(struct msg_t));
			bytes = SSL_read(ssl, &buf, sizeof(struct msg_t));
			if (bytes > 0) {
				short index;
				if (bytes != sizeof(struct msg_t)) {
					fprintf( stderr,
						 "Received non-standard data from server!\n" );
					//conntinue;
					return 1;
				}
				if (buf.chunk.id == 0) {
					if ((index = get_job_slot()) == FAIL) {
						buf.chunk.id = -1; /* ID -1 means reject (full) */
						sprintf((char*)buf.chunk.data, "The agent's queue is full!");
						SSL_write(ssl, &buf, sizeof(struct msg_t));
						continue;
					}
					args[index].slot = FULL;
					memcpy(&args[index].buf, &buf, sizeof(struct msg_t));
					switch (args[index].buf.meta.type) {
					case UNIX:
						pthread_create( &job_thread[index],
								NULL,
								exec_unix,
								&args[index] );
						continue;
					case INSTALL_PKG:
						pthread_create( &job_thread[index],
								NULL,
								install_pkg,
								&args[index] );
						continue;
					case QUERY_PKG:
						pthread_create( &job_thread[index],
								NULL,
								query_pkg,
								&args[index] );
						continue;
					case DELETE_PKG:
						pthread_create( &job_thread[index],
								NULL,
								delete_pkg,
								&args[index] );
						continue;
					case LIST_PKGS:
						pthread_create( &job_thread[index],
								NULL,
								list_pkgs,
								&args[index] );
						continue;
					case UPDATE_PKG:
						pthread_create( &job_thread[index],
								NULL,
								update_pkg,
								&args[index] );
						continue;
					case UPDATE_PKGS:
						pthread_create( &job_thread[index],
								NULL,
								update_pkgs,
								&args[index] );
						continue;
					case GET_OS:
						pthread_create( &job_thread[index],
								NULL,
								get_os,
								&args[index] );
						continue;
					case GET_KERNEL:
						pthread_create( &job_thread[index],
								NULL,
								get_kernel,
								&args[index] );
						continue;
					case GET_UPTIME:
						pthread_create( &job_thread[index],
								NULL,
								get_uptime,
								&args[index] );
						continue;
					case GET_MEMORY:
						pthread_create( &job_thread[index],
								NULL,
								get_memory,
								&args[index] );
						continue;
					default:
						buf.chunk.id = -1;
						sprintf( (char*)buf.chunk.data,
							 "Unsupported job type with ID: %d",
							 buf.meta.type );
						SSL_write(ssl, &buf, sizeof(struct msg_t));
						continue;
					}
				} else {
					index = find_job(buf.meta.id);
					if (index == FAIL) {
						sprintf( (char*)buf.chunk.data,
							 "Data was sent for an invalid job ID" );
						SSL_write(ssl, &buf, sizeof(struct msg_t));
					} else
						memcpy(&args[index].buf, &buf, sizeof(struct msg_t));
				}
			}
		} while (bytes);
		if (SSL_get_shutdown(ssl)) {
			printf("RMPS server has shutdown, trying to reconnect...\n");
			sleep(5);
		} else {
			err = SSL_get_error(ssl, bytes);
				switch (err) {
				case SSL_ERROR_WANT_WRITE:
				printf("want_write\n");
					return 0;
				case SSL_ERROR_WANT_READ:
				printf("want_read\n");
					return 0;
				case SSL_ERROR_ZERO_RETURN:
				printf("zero_return\n");
					return -1;
				case SSL_ERROR_SYSCALL:
				printf("syscall\n");
					return -1;
				case SSL_ERROR_SSL:
				printf("ssl\n");
					return -1;
				default:
					return -1;
			}
		}
	}
}