#!/usr/bin/python3
#
#	toyecc - A small Elliptic Curve Cryptography Demonstration.
#	Copyright (C) 2011-2022 Johannes Bauer
#
#	This file is part of toyecc.
#
#	toyecc 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; this program is ONLY licensed under
#	version 3 of the License, later versions are explicitly excluded.
#
#	toyecc 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 toyecc; if not, write to the Free Software
#	Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
#	Johannes Bauer <JohannesBauer@gmx.de>
#

import struct
import base64
import sys
import toyecc
import hashlib
from toyecc.Random import secure_rand

from MultiCommand import MultiCommand

Ed25519 = toyecc.getcurvebyname("ed25519")

class SignifySignature(object):
	_STRUCT = struct.Struct("< 2s 8s 64s")

	def __init__(self, signature, fingerprint):
		self._signature = signature
		self._fingerprint = fingerprint

	@property
	def fingerprint(self):
		return self._fingerprint

	@property
	def signature(self):
		return self._signature

	@staticmethod
	def fromfile(filename):
		f = open(filename, "r")
		f.readline()
		line = f.readline()
		data = base64.b64decode(line.strip())

		(header, fingerprint, signature) = SignifySignature._STRUCT.unpack(data)
		assert(header == b"Ed")
		f.close()

		signature = toyecc.ECPrivateKey.EDDSASignature.decode(Ed25519, signature)
		return SignifySignature(signature, fingerprint)

	def __str__(self):
		return "Signature<Fingerprint = %s, Signature = %s>" % (self.fingerprint, self.signature)

class SignifyPublicKey(object):
	_STRUCT = struct.Struct("< 2s 8s 32s")

	def __init__(self, pubkey, fingerprint):
		self._pubkey = pubkey
		self._fingerprint = fingerprint

	@property
	def fingerprint(self):
		return self._fingerprint

	@property
	def pubkey(self):
		return self._pubkey

	@staticmethod
	def fromfile(filename):
		f = open(filename, "r")
		f.readline()
		line = f.readline()
		data = base64.b64decode(line.strip())

		(header, fingerprint, pubkey) = SignifyPublicKey._STRUCT.unpack(data)
		assert(header == b"Ed")
		f.close()

		pubkey = toyecc.ECPublicKey.eddsa_decode(Ed25519, pubkey)
		return SignifyPublicKey(pubkey, fingerprint)

	def writetofile(self, filename):
		header = b"Ed"
		data = self._STRUCT.pack(header, self._fingerprint, self._pubkey.eddsa_encode())

		f = open(filename, "w")
		print("untrusted comment: pysignify public key", file = f)
		print(base64.b64encode(data).decode("utf-8"), file = f)
		f.close()

	def __str__(self):
		return "PubKey<Fingerprint = %s, PubKey = %s>" % (self.fingerprint, self.pubkey)


class SignifyPrivateKey(object):
	_STRUCT = struct.Struct("< 2s 2s I 16s 8s 8s 64s")

	def __init__(self, privkey, fingerprint):
		self._privkey = privkey
		self._fingerprint = fingerprint

	@property
	def fingerprint(self):
		return self._fingerprint

	@property
	def privkey(self):
		return self._privkey

	@staticmethod
	def generate():
		privkey = toyecc.ECPrivateKey.eddsa_generate(Ed25519)
		fingerprint = secure_rand(8)
		return SignifyPrivateKey(privkey, fingerprint)

	@staticmethod
	def fromfile(filename):
		f = open(filename, "r")
		f.readline()
		line = f.readline()
		data = base64.b64decode(line.strip())

		(header, kdf, kdfrounds, salt, checksum, fingerprint, seckey) = SignifyPrivateKey._STRUCT.unpack(data)
		assert(header == b"Ed")
		assert(kdf == b"BK")
		f.close()

		if kdfrounds > 0:
			raise Exception("Encrypted private keys are unsupported at the moment.")

		assert(hashlib.sha512(seckey).digest()[:8] == checksum)

		privkey = toyecc.ECPrivateKey.eddsa_decode(Ed25519, seckey[:32])
		return SignifyPrivateKey(privkey, fingerprint)

	def getpubkey(self):
		return SignifyPublicKey(self._privkey.pubkey, self._fingerprint)

	def writetofile(self, filename):
		header = b"Ed"
		kdf = b"BK"
		kdfrounds = 0
		salt = bytes(16)
		seckey = self._privkey.eddsa_encode() + self._privkey.pubkey.eddsa_encode()
		checksum = hashlib.sha512(seckey).digest()[:8]
		data = self._STRUCT.pack(header, kdf, kdfrounds, salt, checksum, self.fingerprint, seckey)

		f = open(filename, "w")
		print("untrusted comment: pysignify secret key", file = f)
		print(base64.b64encode(data).decode("utf-8"), file = f)
		f.close()

	def __str__(self):
		return "PrivKey<Fingerprint = %s, PubKey = %s>" % (self.fingerprint, self.pubkey)


def action_verify(cmd, args):
	if (args.contentfile is None) and (args.sigfile.endswith(".sig")):
		contentfile = args.sigfile[:-4]
	else:
		contentfile = args.contentfile
	if contentfile is None:
		raise Exception("Cannot determine content file automatically, please specify.")

	pubkey = SignifyPublicKey.fromfile(args.pubkeyfile)
	if args.verbose:
		print("Public key:", pubkey)
	signature = SignifySignature.fromfile(args.sigfile)
	content = open(contentfile, "rb").read()

	if pubkey.fingerprint != signature.fingerprint:
		print("Warning: Fingerprint of supplied public key and signature file do not match.")

	if pubkey.pubkey.eddsa_verify(content, signature.signature):
		print("Signature verification successful.")
		sys.exit(0)
	else:
		print("Signature verification FAILED!")
		sys.exit(1)

def action_sign(cmd, args):
	if args.sigfile is None:
		sigfile = args.contentfile + ".sig"
	else:
		sigfile = args.sigfile

	privkey = SignifyPrivateKey.fromfile(args.privkeyfile)
	content = open(args.contentfile, "rb").read()
	signature = privkey.privkey.eddsa_sign(content)

	signed_data = bytearray()
	signed_data += b"Ed"
	signed_data += privkey.fingerprint
	signed_data += signature.encode()

	f = open(sigfile, "w")
	print("untrusted comment: signed by pysignify", file = f)
	print(base64.b64encode(signed_data).decode("utf-8"), file = f)
	f.close()

def action_genkey(cmd, args):
	privkey = SignifyPrivateKey.generate()
	privkey.writetofile(args.privkeyfile)
	privkey.getpubkey().writetofile(args.pubkeyfile)

def action_getpubkey(cmd, args):
	privkey = SignifyPrivateKey.fromfile(args.privkeyfile)
	privkey.getpubkey().writetofile(args.pubkeyfile)

mc = MultiCommand(help = "pysignify: Python example to create EdDSA signatures on basis of Ed25519 which are used for the 'signify' signatures which are popular in the OpenBSD world. It should be fully compatible to the signature format used by the BSD signify.")

def genparser(parser):
	parser.add_argument("-p", "--pubkeyfile", metavar = "filename", type = str, required = True, help = "File in which the public key is stored. Mandatory argument.")
	parser.add_argument("-s", "--sigfile", metavar = "filename", type = str, required = True, help = "File in which the signature is stored. Mandatory argument.")
	parser.add_argument("-c", "--contentfile", metavar = "filename", type = str, help = "File in which the content is stored. Will be the sigfile without trailing extension if omitted.")
	parser.add_argument("--verbose", action = "store_true", help = "Increase verbosity during the verification process.")
mc.register("verify", "Verify a signify signature", genparser, action = action_verify)

def genparser(parser):
	parser.add_argument("-k", "--privkeyfile", metavar = "filename", type = str, required = True, help = "File in which the private key is stored. Mandatory argument.")
	parser.add_argument("-c", "--contentfile", metavar = "filename", type = str, required = True, help = "File in which the content that is to be signed is stored.  Mandatory argument.")
	parser.add_argument("-s", "--sigfile", metavar = "filename", type = str, help = "File in which the signature is stored. Will be the content file plus '.sig' extension if omitted.")
	parser.add_argument("--verbose", action = "store_true", help = "Increase verbosity during the signing process.")
mc.register("sign", "Sign a document", genparser, action = action_sign)

def genparser(parser):
	parser.add_argument("-k", "--privkeyfile", metavar = "filename", type = str, required = True, help = "File in which the private key is stored. Mandatory argument.")
	parser.add_argument("-p", "--pubkeyfile", metavar = "filename", type = str, required = True, help = "File in which the public key is stored. Mandatory argument.")
	parser.add_argument("--verbose", action = "store_true", help = "Increase verbosity during the key generation process.")
mc.register("genkey", "Generate a private/public keypair", genparser, action = action_genkey)

def genparser(parser):
	parser.add_argument("-k", "--privkeyfile", metavar = "filename", type = str, required = True, help = "File in which the private key is stored. Mandatory argument.")
	parser.add_argument("-p", "--pubkeyfile", metavar = "filename", type = str, required = True, help = "File in which the public key shall be stored. Mandatory argument.")
	parser.add_argument("--verbose", action = "store_true", help = "Increase verbosity during the public key extraction.")
mc.register("getpubkey", "Extract the public key from the private key file", genparser, action = action_getpubkey)

mc.run(sys.argv[1:])


