#!/usr/bin/env python
#
# mmgen = Multi-Mode GENerator, command-line Bitcoin cold storage solution
# Copyright (C)2013-2017 Philemon <mmgen-py@yandex.com>
#
# 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 3 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/>.

"""
txsign: Sign a transaction generated by 'mmgen-txcreate'
"""

from mmgen.common import *
from mmgen.seed import *
from mmgen.tx import *
from mmgen.addr import *

pnm = g.proj_name

txsign_notes = """
Transactions may contain both {pnm} or non-{pnm} input addresses.

To sign non-{pnm} inputs, a bitcoind wallet dump or flat key list is used
as the key source ('--keys-from-file' option).

To sign {pnm} inputs, key data is generated from a seed as with the
{pnl}-addrgen and {pnl}-keygen commands.  Alternatively, a key-address file
may be used (--mmgen-keys-from-file option).

Multiple wallets or other seed files can be listed on the command line in
any order.  If the seeds required to sign the transaction's inputs are not
found in these files (or in the default wallet), the user will be prompted
for seed data interactively.

To prevent an attacker from crafting transactions with bogus {pnm}-to-Bitcoin
address mappings, all outputs to {pnm} addresses are verified with a seed
source.  Therefore, seed files or a key-address file for all {pnm} outputs
must also be supplied on the command line if the data can't be found in the
default wallet.

Seed source files must have the canonical extensions listed in the 'FileExt'
column below:

  {f}
""".format(f='\n  '.join(SeedSource.format_fmt_codes().splitlines()),
			pnm=pnm,pnl=pnm.lower())

wmsg = {
	'mapping_error': """
{pnm} -> BTC address mappings differ!
%-23s %s -> %s
%-23s %s -> %s
""".strip().format(pnm=pnm),
	'missing_keys_error': """
ERROR: a key file must be supplied for the following non-{pnm} address%s:\n    %s
""".format(pnm=pnm).strip()
}

saved_seeds = {}

def get_seed_for_seed_id(sid,infiles,saved_seeds):

	if sid in saved_seeds:
		return saved_seeds[sid]

	while True:
		if infiles:
			ss = SeedSource(infiles.pop(0),ignore_in_fmt=True)
		elif opt.in_fmt:
			qmsg('Need seed data for Seed ID %s' % sid)
			ss = SeedSource()
			msg('User input produced Seed ID %s' % ss.seed.sid)
		else:
			die(2,'ERROR: No seed source found for Seed ID: %s' % sid)

		saved_seeds[ss.seed.sid] = ss.seed
		if ss.seed.sid == sid: return ss.seed

def generate_keys_for_mmgen_addrs(mmgen_addrs,infiles,saved_seeds):
	sids = set(i.sid for i in mmgen_addrs)
	vmsg('Need seed%s: %s' % (suf(sids,'s'),' '.join(sids)))
	d = AddrListList()
	from mmgen.addr import KeyAddrList
	for sid in sids:
		# Returns only if seed is found
		seed = get_seed_for_seed_id(sid,infiles,saved_seeds)
		for t in MMGenAddrType.mmtypes:
			idx_list = [i.idx for i in mmgen_addrs if i.sid == sid and i.mmtype == t]
			if idx_list:
				addr_idxs = AddrIdxList(idx_list=idx_list)
				d += KeyAddrList(seed=seed,addr_idxs=addr_idxs,do_chksum=False,mmtype=MMGenAddrType(t)).flat_list()
	return d

def add_keys(tx,src,infiles=None,saved_seeds=None,keyaddr_list=None):
	need_keys = [e for e in getattr(tx,src) if e.mmid and not e.have_wif]
	if not need_keys: return []
	desc,m1 = ('key-address file','From key-address file:') if keyaddr_list else \
					('seed(s)','Generated from seed:')
	qmsg('Checking {} -> BTC address mappings for {} (from {})'.format(pnm,src,desc))
	d = keyaddr_list.flat_list() if keyaddr_list else \
		generate_keys_for_mmgen_addrs([e.mmid for e in need_keys],infiles,saved_seeds)
	new_keys = []
	for e in need_keys:
		for f in d:
			if f.mmid == e.mmid:
				if f.addr == e.addr:
					e.have_wif = True
					if src == 'inputs':
						new_keys.append((f.addr,f.wif))
				else:
					die(3,wmsg['mapping_error'] % (m1,f.mmid,f.addr,'tx file:',e.mmid,e.addr))
	if new_keys:
		vmsg('Added %s wif key%s from %s' % (len(new_keys),suf(new_keys,'s'),desc))
	return new_keys

def _pop_and_return(args,cmplist): # strips found args
	return list(reversed([args.pop(args.index(a)) for a in reversed(args) if get_extension(a) in cmplist]))

def get_tx_files(opt,args):
	ret = _pop_and_return(args,[MMGenTX.raw_ext])
	if not ret: die(1,'You must specify a raw transaction file!')
	return ret

def get_seed_files(opt,args):
	# favor unencrypted seed sources first, as they don't require passwords
	u,e = SeedSourceUnenc,SeedSourceEnc
	ret = _pop_and_return(args,u.get_extensions())
	from mmgen.filename import find_file_in_dir
	wf = find_file_in_dir(Wallet,g.data_dir) # Make this the first encrypted ss in the list
	if wf: ret.append(wf)
	ret += _pop_and_return(args,e.get_extensions())
	if not (ret or opt.mmgen_keys_from_file or opt.keys_from_file): # or opt.use_wallet_dat
		die(1,'You must specify a seed or key source!')
	return ret

def get_keyaddrlist(opt):
	if opt.mmgen_keys_from_file:
		return KeyAddrList(opt.mmgen_keys_from_file)
	return None

def get_keylist(opt):
	if opt.keys_from_file:
		l = get_lines_from_file(opt.keys_from_file,'key-address data',trim_comments=True)
		ret = KeyAddrList(keylist=[m.split()[0] for m in l]) # accept bitcoind wallet dumps
		ret.generate_addrs_from_keylist()
		return ret
	return None

def txsign(opt,c,tx,seed_files,kl,kal,tx_num_str=''):
	# Start
	keys = []
#	tx.pmsg()
	non_mm_addrs = tx.get_non_mmaddrs('inputs')
	if non_mm_addrs:
		tmp = KeyAddrList(addrlist=non_mm_addrs,do_chksum=False)
		tmp.add_wifs(kl)
		m = tmp.list_missing('wif')
		if m: die(2,wmsg['missing_keys_error'] % (suf(m,'es'),'\n    '.join(m)))
		keys += tmp.get_addr_wif_pairs()

	if opt.mmgen_keys_from_file:
		keys += add_keys(tx,'inputs',keyaddr_list=kal)
		add_keys(tx,'outputs',keyaddr_list=kal)

	keys += add_keys(tx,'inputs',seed_files,saved_seeds)
	add_keys(tx,'outputs',seed_files,saved_seeds)

	tx.delete_attrs('inputs','have_wif')
	tx.delete_attrs('outputs','have_wif')

	extra_sids = set(saved_seeds) - tx.get_input_sids() - tx.get_output_sids()
	if extra_sids:
		msg('Unused Seed ID{}: {}'.format(suf(extra_sids,'s'),' '.join(extra_sids)))

	if tx.sign(c,tx_num_str,dict(keys)):
		return tx
	else:
		die(3,'failed\nSome keys were missing.  Transaction {}could not be signed.'.format(tx_num_str))