#!/usr/bin/env python3

import sys
import json
import os
import uuid
from datetime import datetime

import scrypt
from ecdsa import SECP256k1, SigningKey
from Crypto.Cipher import AES
from Crypto.Util import Counter
from Crypto.Hash import keccak

scryptN = 1 << 18
scryptP = 1
scryptR = 8
scryptLen = 32
aesBlockSize = 1 << 7
pkLen = 32
ivLen = 16
saltLen = 32
pwLen = 64
walletVersion = 3

# uuid for the wallet
uu = uuid.uuid1()

# make new wallet example
iv = os.urandom(ivLen)
salt = os.urandom(saltLen)
pw = os.urandom(pwLen)

# save the password to file
fd = os.open(".swarm-pass", os.O_WRONLY | os.O_CREAT, mode=0o600)
c = os.write(fd, pw)
if (c != pwLen):
	raise Exception("Incomplete password write to file")
os.close(fd)

# derive the aes key used to encrypt the private key
key = scrypt.hash(pw, salt, scryptN, scryptR, scryptP, scryptLen)

# generate the keypair
pk = SigningKey.generate(curve=SECP256k1)
pkBytes = pk.to_string()
pubBytes = pk.get_verifying_key().to_string()

# hash and truncate public key to get the ethereum address
hasher = keccak.new(digest_bits=256)
hasher.update(pubBytes)
addrBytes = hasher.digest()[12:]
addrHex = addrBytes.hex()
#print("{}\n{}\n{}\n".format(pkBytes.hex(), pubBytes.hex(), addrHex))

# encrypt the private key with the key generated by the password
ctr = Counter.new(aesBlockSize, initial_value=int.from_bytes(iv, "big"))
cipher = AES.new(key[:16], AES.MODE_CTR, counter=ctr)
cipherText = cipher.encrypt(pkBytes)

# calculate the mac
hasher = keccak.new(digest_bits=256)
hasher.update(key[16:])
hasher.update(cipherText)
mac = hasher.digest()

# format to json wallet
wallet = """{
  "address": \"""" + addrHex + """\",
  "crypto": {
    "cipher": "aes-128-ctr",
    "ciphertext": \"""" + cipherText.hex() + """\",
    "cipherparams": {
      "iv": \"""" + iv.hex() + """\"
    },
    "kdf": "scrypt",
    "kdfparams": {
      "dklen": """ + str(scryptLen) + """,
      "n": """ + str(scryptN) + """,
      "p": """ + str(scryptP) + """,
      "r": """ + str(scryptR) + """,
      "salt": \"""" + salt.hex() + """\"
    },
    "mac": \"""" + mac.hex() + """\"
  },
  "id": \"""" + str(uu) + """\",
  "version": """ + str(walletVersion) + """
}"""

# write wallet to file with standardized filename
now = datetime.now()
walletFilename = now.strftime("UTC--%Y-%m-%dT%H-%M-%S.%f000Z--" + addrHex)
fd = os.open(walletFilename, os.O_WRONLY | os.O_CREAT, 0o600)
c = os.write(fd, wallet.encode('utf-8'))
os.close(fd)