/* Cryptographic card PIN check and RSA decryption utility
 * Copyright (C) 2015 - 2020 Timothy Pearson <kb9vqf@pearsoncomputing.net>
 *
 * 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 program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <stdint.h>
#include <termios.h>
#include <unistd.h>
#include <string.h>

#include <pkcs11-helper-1.0/pkcs11h-certificate.h>
#include <pkcs11-helper-1.0/pkcs11h-openssl.h>
#include <openssl/x509v3.h>

#define CARD_MAX_LOGIN_RETRY_COUNT 3

char has_plymouth = 0;
char use_cached_pin = 0;
char* cached_pin = NULL;

static PKCS11H_BOOL pkcs_pin_hook(IN void * const global_data, IN void * const user_data, IN const pkcs11h_token_id_t token, IN const unsigned retry, OUT char * const pin, IN const size_t pin_max) {
	int pos;
	char *line = NULL;
	size_t size;
	ssize_t read;

	if (use_cached_pin && cached_pin) {
		// Copy PIN to buffer
		snprintf(pin, pin_max, "%s", cached_pin);

		// Success
		return 1;
	}

	// Hide input
	struct termios oldt;
	tcgetattr(STDIN_FILENO, &oldt);
	struct termios newt = oldt;
	newt.c_lflag &= ~ECHO;
	tcsetattr(STDIN_FILENO, TCSANOW, &newt);

	if (has_plymouth) {
		char buffer[1024];
		snprintf(buffer, 1024, "plymouth ask-for-password --prompt=\"Please enter the PIN for '%s'\"", token->display);
		system(buffer);
	}
	else {
		fprintf(stderr, "Please enter the PIN for '%s'\n", token->display);
	}
	fflush(stdout);

	read = getline(&line, &size, stdin);
	if ((read < 0) || (read >= pin_max)) {
		free(line);

		// Abort
		tcsetattr(STDIN_FILENO, TCSANOW, &oldt);
		return 0;
	}
	else {
		// Strip newlines
		pos = 0;
		while (line[pos] != 0) {
			if ((line[pos] == '\n') || (line[pos] == '\r')) {
				line[pos] = 0;
				break;
			}
			pos++;
		}

		// Copy PIN to cache
		if (cached_pin) {
			free(cached_pin);
		}
		cached_pin = malloc(sizeof(char) * pin_max);
		snprintf(cached_pin, pin_max, "%s", line);

		// Copy PIN to buffer
		snprintf(pin, pin_max, "%s", line);
		free(line);

		// Success
		tcsetattr(STDIN_FILENO, TCSANOW, &oldt);
		return 1;
	}
}

static void pkcs_log_hook(IN void * const global_data, IN unsigned flags, IN const char * const format, IN va_list args) {
	if (!has_plymouth) {
		vfprintf(stderr, format, args);
		fprintf(stderr, "\n");
	}
}

char* tde_autopin(X509* x509_cert) {
	char* retString = NULL;
	int i;

	// Use subjAltName field in card certificate to provide the card's PIN,
	// in order to support optional pin-less operation.
	// Parse the TDE autologin extension
	// Structure:
	// OID 1.3.6.1.4.1.40364.1.2.1
	// 	SEQUENCE
	// 		ASN1_CONSTRUCTED [index: 0] (field name: pin)
	// 			GeneralString

	// Register custom OID type for TDE autopin data
	ASN1_OBJECT* tde_autopin_data_object = OBJ_txt2obj("1.3.6.1.4.1.40364.1.2.1", 0);

	GENERAL_NAMES* subjectAltNames = (GENERAL_NAMES*)X509_get_ext_d2i(x509_cert, NID_subject_alt_name, NULL, NULL);
	int altNameCount = sk_GENERAL_NAME_num(subjectAltNames);
	for (i=0; i < altNameCount; i++) {
		GENERAL_NAME* generalName = sk_GENERAL_NAME_value(subjectAltNames, i);
		if (generalName->type == GEN_OTHERNAME) {
			OTHERNAME* otherName = generalName->d.otherName;
			if (!OBJ_cmp(otherName->type_id, tde_autopin_data_object)) {
				ASN1_TYPE* asnValue = otherName->value;
				if (asnValue) {
					// Found autopin structure
					int index;
					ASN1_TYPE* asnSeqValue = NULL;
					ASN1_GENERALSTRING* asnGeneralString = NULL;
					STACK_OF(ASN1_TYPE) *asnSeqValueStack = NULL;
					long asn1SeqValueObjectLength;
					int asn1SeqValueObjectTag;
					int asn1SeqValueObjectClass;
					int returnCode;
					const uint8_t* asnSeqValueString = ASN1_STRING_get0_data(asnValue->value.sequence);

					index = 0;	// Search for the PIN field
					asnSeqValueStack = d2i_ASN1_SEQUENCE_ANY(NULL, &asnSeqValueString, ASN1_STRING_length(asnValue->value.sequence));
					asnSeqValue = sk_ASN1_TYPE_value(asnSeqValueStack, index);
					if (asnSeqValue) {
						if (asnSeqValue->value.octet_string->data[0] == ((V_ASN1_CONSTRUCTED | V_ASN1_CONTEXT_SPECIFIC) + index)) {
							const unsigned char* asn1SeqValueObjectData = asnSeqValue->value.sequence->data;
							returnCode = ASN1_get_object(&asn1SeqValueObjectData, &asn1SeqValueObjectLength, &asn1SeqValueObjectTag, &asn1SeqValueObjectClass, asnSeqValue->value.sequence->length);
							if (!(returnCode & 0x80)) {
								if (returnCode == (V_ASN1_CONSTRUCTED + index)) {
									if (d2i_ASN1_GENERALSTRING(&asnGeneralString, &asn1SeqValueObjectData, asn1SeqValueObjectLength) != NULL) {
										retString = strdup(ASN1_STRING_get0_data(asnGeneralString));
									}
								}
							}
						}
					}
					sk_ASN1_TYPE_pop_free(asnSeqValueStack, ASN1_TYPE_free);
				}
			}
		}
	}

	// Clean up
	OBJ_cleanup();

	return retString;
}

int main(int argc, char* argv[]) {
	CK_RV rv;
	pkcs11h_certificate_id_list_t issuers;
	pkcs11h_certificate_id_list_t certs;

	has_plymouth = 0;
	const char* with_plymount_var = getenv("HAS_PLYMOUTH");
	if (with_plymount_var && (with_plymount_var[0] == '1')) {
		has_plymouth = 1;
	}

	if ((argc < 2) || (argv[1][0] == 0)) {
		fprintf(stderr, "Usage: ./cardpincheck <opensc provider library> <file to decrypt>\n"
			"Example: ./cardpincheck /usr/lib/opensc-pkcs11.so\n");
		return -5;
	}

	char* opensc_provider_library = argv[1];

	char decryption_requested = 0;
	char* file_to_decrypt = NULL;
	if (argc > 2) {
		decryption_requested = 1;
		file_to_decrypt = argv[2];
	}

	fprintf(stderr, "Initializing pkcs11-helper\n"); fflush(stderr);
	if ((rv = pkcs11h_initialize()) != CKR_OK) {
		fprintf(stderr, "pkcs11h_initialize failed: %s\n", pkcs11h_getMessage(rv));
		return -1;
	}

	fprintf(stderr, "Registering pkcs11-helper hooks\n"); fflush(stderr);
	if ((rv = pkcs11h_setLogHook(pkcs_log_hook, NULL)) != CKR_OK) {
		fprintf(stderr, "pkcs11h_setLogHook failed: %s\n", pkcs11h_getMessage(rv));
		return -1;
	}
	pkcs11h_setLogLevel(PKCS11H_LOG_WARN);
	// pkcs11h_setLogLevel(PKCS11H_LOG_DEBUG2);

#if 0
	if ((rv = pkcs11h_setTokenPromptHook(_pkcs11h_hooks_token_prompt, NULL)) != CKR_OK) {
		fprintf(stderr, "pkcs11h_setTokenPromptHook failed: %s\n", pkcs11h_getMessage(rv));
		return -1;
	}
#endif

	if ((rv = pkcs11h_setMaxLoginRetries(CARD_MAX_LOGIN_RETRY_COUNT)) != CKR_OK) {
		fprintf(stderr, "pkcs11h_setMaxLoginRetries failed: %s\n", pkcs11h_getMessage(rv));
		return -1;
	}

	if ((rv = pkcs11h_setPINPromptHook(pkcs_pin_hook, NULL)) != CKR_OK) {
		fprintf(stderr, "pkcs11h_setPINPromptHook failed: %s\n", pkcs11h_getMessage(rv));
		return -1;
	}

	fprintf(stderr, "Adding provider '%s'\n", opensc_provider_library); fflush(stderr);
		if ((rv = pkcs11h_addProvider(opensc_provider_library, opensc_provider_library, FALSE, PKCS11H_PRIVATEMODE_MASK_AUTO, PKCS11H_SLOTEVENT_METHOD_AUTO, 0, FALSE)) != CKR_OK) {
		fprintf(stderr, "pkcs11h_addProvider failed: %s\n", pkcs11h_getMessage(rv));
		return -1;
	}

	rv = pkcs11h_certificate_enumCertificateIds(PKCS11H_ENUM_METHOD_CACHE, NULL, PKCS11H_PROMPT_MASK_ALLOW_PIN_PROMPT, &issuers, &certs);
	if ((rv != CKR_OK) || (certs == NULL)) {
		fprintf(stderr, "Cannot enumerate certificates: %s\n", pkcs11h_getMessage(rv));
		return -1;
	}

	int ret = -1;
	int i = 0;
	pkcs11h_certificate_id_list_t cert;
	pkcs11h_certificate_t certificate = NULL;
	RSA* rsa_pubkey = NULL;
	for (cert = certs; cert != NULL; cert = cert->next) {
		rv = pkcs11h_certificate_create(certs->certificate_id, NULL, PKCS11H_PROMPT_MASK_ALLOW_PIN_PROMPT, PKCS11H_PIN_CACHE_INFINITE, &certificate);
		if (rv != CKR_OK) {
			fprintf(stderr, "Cannot read certificate: %s\n", pkcs11h_getMessage(rv));
			pkcs11h_certificate_freeCertificateId(certs->certificate_id);
			ret = -1;
			break;
		}

		pkcs11h_certificate_freeCertificateId(certs->certificate_id);

		pkcs11h_openssl_session_t openssl_session = NULL;
		if ((openssl_session = pkcs11h_openssl_createSession(certificate)) == NULL) {
			fprintf(stderr, "Cannot initialize openssl session to retrieve cryptographic objects\n");
			pkcs11h_certificate_freeCertificate(certificate);
			ret = -1;
			break;
		}

		// Get certificate data
		X509* x509_local;
		x509_local = pkcs11h_openssl_session_getX509(openssl_session);
		if (!x509_local) {
			fprintf(stderr, "Cannot get X509 object\n");
			ret = -1;
		}

		// Check for TDE autopin structure
		char* autopin = tde_autopin(x509_local);
		if (autopin) {
			cached_pin = autopin;
			use_cached_pin = 1;
		}

		// Extract public key from X509 certificate
		EVP_PKEY* x509_pubkey = NULL;
		x509_pubkey = X509_get_pubkey(x509_local);
		if (x509_pubkey) {
			rsa_pubkey = EVP_PKEY_get1_RSA(x509_pubkey);
		}

		// Check PIN
		rv = pkcs11h_certificate_ensureKeyAccess(certificate);
		if (rv != CKR_OK) {
			if (rv == CKR_GENERAL_ERROR) {
				ret = -4;
				break;
			}
			else if (rv == CKR_CANCEL) {
				ret = -3;
				break;
			}
			else if ((rv == CKR_PIN_INCORRECT) || (rv == CKR_USER_NOT_LOGGED_IN)) {
				ret = -2;
				break;
			}
			else {
				ret = -2;
				break;
			}
		}
		else {
			// Success!
			ret = 0;
			break;
		}

		pkcs11h_certificate_freeCertificate(certificate);
		certificate = NULL;

		i++;
	}

	if (decryption_requested && (ret == 0)) {
		// We know the cached PIN is correct; disable any further login prompts
		use_cached_pin = 1;

		char abort_decryption = 0;
		if (file_to_decrypt) {
			long ciphertextfilesize = 0;
			FILE *ciphertextfile = fopen(file_to_decrypt, "r");
			if (ciphertextfile) {
				fseek(ciphertextfile, 0, SEEK_END);
				ciphertextfilesize = ftell(ciphertextfile);
				fseek(ciphertextfile, 0, SEEK_SET);

				char* ciphertext = malloc(ciphertextfilesize + 1);
				fread(ciphertext, ciphertextfilesize, 1, ciphertextfile);
				fclose(ciphertextfile);

				// Verify minimum size
				if (ciphertextfilesize < 16) {
					fprintf(stderr, "Cannot decrypt: ciphertext too small\n");
					abort_decryption = 1;
				}

				// Try to get RSA parameters and verify maximum size
				if (rsa_pubkey) {
					unsigned int rsa_length = RSA_size(rsa_pubkey);
					if (ciphertextfilesize > rsa_length) {
						fprintf(stderr, "Cannot decrypt: ciphertext too large\n");
						abort_decryption = 1;
					}
				}

				if (!abort_decryption) {
					// Try decryption
					size_t size = 0;

					// Determine output buffer size
					rv = pkcs11h_certificate_decryptAny(certificate, CKM_RSA_PKCS, ciphertext, ciphertextfilesize, NULL, &size);
					if (rv != CKR_OK) {
						fprintf(stderr, "Cannot determine decrypted message length: %s (%d)\n", pkcs11h_getMessage(rv), rv);
						if (rv == CKR_FUNCTION_FAILED) {
							/* Decryption failed */
							ret = -20;
							abort_decryption = 1;
						}
						else if (rv == CKR_CANCEL) {
							ret = -2;
							abort_decryption = 1;
						}
						else if ((rv == CKR_PIN_INCORRECT) || (rv == CKR_USER_NOT_LOGGED_IN)) {
							ret = -1;
							abort_decryption = 1;
						}
						else {
							abort_decryption = 1;
						}
					}
					else {
						// Decrypt data
						char* plaintext = malloc(size);
						rv = pkcs11h_certificate_decryptAny(certificate, CKM_RSA_PKCS, ciphertext, ciphertextfilesize, plaintext, &size);
						if (rv != CKR_OK) {
							fprintf(stderr, "Cannot decrypt: %s (%d)\n", pkcs11h_getMessage(rv), rv);
							if (rv == CKR_CANCEL) {
								ret = -1;
								abort_decryption = 1;
							}
							else if ((rv == CKR_PIN_INCORRECT) || (rv == CKR_USER_NOT_LOGGED_IN)) {
								ret = -1;
								abort_decryption = 1;
							}
						}
						else {
							// Write decrypted data to stdout
							fwrite(plaintext, sizeof(char), size, stdout);
							fflush(stdout);
						}
						free(plaintext);
					}
				}
				free(ciphertext);
			}
		}
	}
	else if (ret == 0) {
		printf("%s", cached_pin);
	}

	if (certificate) {
		pkcs11h_certificate_freeCertificate(certificate);
	}
	pkcs11h_certificate_freeCertificateIdList(issuers);

	return ret;
}