AsyncWalletGeneration.scala

package co.upvest.martind

import co.upvest.terminology.adjectives.common._
import co.upvest.terminology.adjectives.implicits._

import co.upvest.dry.essentials._
import co.upvest.dry.essentials.syntax._
import co.upvest.dry.catz.syntax._
import co.upvest.dry.{gstorage => gs}
import co.upvest.dry.service.{UpvestMsgId, Flogging}
import co.upvest.dry.cryptoadt.bitcoin
import co.upvest.dry.cryptoadt.algorithms.Scrypt
import co.upvest.dry.breakout.api.{UserSecretKey, MasterNodeSeedKey}
import co.upvest.dry.breakout.{Client => BreakoutClient}

import adt._
import algorithms.{HierarchicalDeterministicStretching => HD}

import io.circe
import io.circe.parser._
import io.circe.syntax._

import cats.{~>, MonadError}
import cats.data.EitherT
import cats.instances.future._
import cats.syntax.bifunctor._
import cats.syntax.flatMap._
import cats.instances.either._
import cats.syntax.option._

import scala.concurrent.{ExecutionContext, Future}
import scala.util.Success

class AsyncWalletGeneration(
  gstorage: gs.Client[Future],
  masterNodeSeedBucket: gs.Bucket,
  breakout: BreakoutClient,
  userSecretKey: UserSecretKey,
  masterNodeSeedKey: MasterNodeSeedKey,
  implicit val bitcoinNetwork: bitcoin.Network,
  defaultKdfParams: Scrypt.Params.Overrides,
)(implicit ec: ExecutionContext) extends MetadataInstances {

  implicit def insertDefaultUserSecretKey(us: UserSecret) =
    co.upvest.dry.breakout.api.UserSecret(userSecretKey, us.s);

  def extractKdfParamsOverrides(
    r: { def kdfparams: Option[Scrypt.Params.Overrides] }
  ): Option[Scrypt.Params.Overrides] = (
    r.kdfparams getOrElse Scrypt.Params.Overrides.Empty
      orElse defaultKdfParams some
  )

  sealed trait Failure {
    def message: String
    def throwable: Option[Throwable]
    def upvestMsgId: Option[UpvestMsgId]
  }

  object Failure {
    case class Parsing(t: circe.ParsingFailure) extends Failure {
      lazy val message = "failed to parse object"
      lazy val throwable = Some(t)
      lazy val upvestMsgId = None
    }

    case class Decoding(t: circe.DecodingFailure) extends Failure {
      lazy val message = "failed to decode object"
      lazy val throwable = Some(t)
      lazy val upvestMsgId = None
    }

    case class InvalidStringEncoding(t: Throwable) extends Failure {
      lazy val message = "invalid string encoding"
      lazy val throwable = Some(t)
      lazy val upvestMsgId = None
    }

    case class GStorage(r: Request, t: Throwable) extends Failure {
      lazy val message = "failure in GStorage"
      lazy val throwable = Some(t)
      lazy val upvestMsgId = Some(r.upvestMsgId)
    }

    case class UnknownSeedHash(r: Request, sh: RawSeedHash) extends Failure {
      lazy val message = s"unknown seed hash: ${sh.b64u}"
      lazy val throwable = None
      lazy val upvestMsgId = Some(r.upvestMsgId)
    }

    case class MACMismatch(
      r: Request,
      kind: String,
      mac1: Hashed[Bytes],
      mac2: Hashed[Bytes]
    ) extends Failure {
      lazy val message = s"$kind MAC mismatch: ${mac1.as.base64} =/= ${mac2.as.base64}"
      lazy val throwable = None
      lazy val upvestMsgId = Some(r.upvestMsgId)
    }

    case class SeedHashMismatch(
      r: Request,
      sh1: Hashed[Bytes],
      sh2: Hashed[Bytes]
    ) extends Failure {
      lazy val message = s"Seed hash mismatch: ${sh1.as.base64URL} =/= ${sh2.as.base64URL}"
      lazy val throwable = None
      lazy val upvestMsgId = Some(r.upvestMsgId)
    }

    case class UnexpectedAddress(r: Request, expected: String, got: String) extends Failure {
      lazy val message = s"Unexpected address: $got expected: $expected"
      lazy val throwable = None
      lazy val upvestMsgId = Some(r.upvestMsgId)
    }

    case class Breakout(r: Request, t: Throwable) extends Failure {
      lazy val message = "call to the breakout service failed"
      lazy val throwable = Some(t)
      lazy val upvestMsgId = Some(r.upvestMsgId)
    }
  }

  type F[A] = EitherT[Future, Failure, A]
  private val M = implicitly[MonadError[F, this.Failure]]

  final val SupportedAssetKinds: Set[AssetKind] =
    Set(AssetKind.Arweave, AssetKind.Bitcoin, AssetKind.Ethereum)

  private def liftFuture(e: Throwable => Failure) =
    new (Future ~> F) {
      override def apply[A](fa: Future[A]): F[A] =
        EitherT(fa transform { _.toEither |> Success.apply }) leftMap e
    }

  private def liftClosure(e: Throwable => Failure) =
    new ((() => ?) ~> F) {
      override def apply[A](fa: () => A): F[A] = liftFuture(e)(Future{ fa() })
    }

  private val underlyingFailure = new ((() => ?) ~> F) {
    override def apply[A](fa: () => A): F[A] = EitherT.right(Future{ fa() })
  }

  private val flogger: Flogging[F] = Flogging.typesafe(
    underlyingFailure,
    name = this.getClass.getName
  )

  private def put(r: Request)(bi: gs.Blob.Id, data: Bytes, ct: gs.ContentType) =
    liftFuture { t => Failure.GStorage(r, t) } {
      gstorage.put(bi, data, ct) >>| { _ => () }
    }

  private def get(r: Request)(bi: gs.Blob.Id) =
    liftFuture { t => Failure.GStorage(r, t) } {
      gstorage.get(bi)
    }

  private def latest(r: Request)(prefix: gs.Name): F[Option[gs.Blob[Future]]] =
    liftFuture { t => Failure.GStorage(r, t) } {
      gstorage.latest(masterNodeSeedBucket, Some(prefix))
    }

  def deserialize[T: circe.Decoder](str: String): F[T] = for {
    json <- M.fromEither { parse(str) leftMap Failure.Parsing }
    t <- M.fromEither { json.as[T] leftMap Failure.Decoding }
  } yield t

  private def filename(sh: RawSeedHash, mac: Bytes): gs.Name =
    s"${filePrefix(sh)}/mac/${mac.base64URL}.json"

  private def filePrefix(sh: RawSeedHash): gs.Name = s"seedHash/${sh.b64u}"

  private def liftBreakout(r: Request) =
    liftFuture { t => Failure.Breakout(r, t) }

  private def storeMasterNodeSeed(r: Request)(
    emns: EncryptedMasterNodeSeed
  ): F[Unit] = for {
    _ <- flogger.debug("Storing master-node-seed", r.upvestMsgId, emns.seedHash)
    _ <- put(r)(
      gs.Blob.Id(
        masterNodeSeedBucket,
        name = filename(emns.seedHash, emns.userSecretEncrypted.mac.as[Bytes]),
        generation = None
      ),
      data = emns.asJson.noSpaces.utf8,
      ct = "application/json"
    )
    _ <- flogger.debug("Stored master-node-seed", r.upvestMsgId, emns.seedHash)
  } yield ()

  private def loadMasterNodeSeedBlob(r: Request)(
    sh: RawSeedHash
  ): F[EncryptedMasterNodeSeed] = for {
    b <- latest(r)(filePrefix(sh)) >>= {
      case Some(b) => M pure b
      case _ => M raiseError[gs.Blob[Future]] Failure.UnknownSeedHash(r, sh)
    }

    data <- liftFuture { Failure.GStorage(r, _) } { b.data() }
    str <- liftClosure (Failure.InvalidStringEncoding) { () =>
      data.utf8.get
    }

    emns <- deserialize[EncryptedMasterNodeSeed](str)
  } yield emns

  def generateMasterNodeSeed(
    r: GenerateMasterNodeSeed
  ): F[GeneratedMasterNodeSeed] = for {
    mns <- liftBreakout(r)(breakout.mns.generate(
      r.userSecret, masterNodeSeedKey.pubkey, r.upvestMsgId,
      extractKdfParamsOverrides(r)
    ))
    _ <- storeMasterNodeSeed(r)(mns)
    _ <- flogger.info("Generated master-node-seed", r.userId, mns, r.upvestMsgId)
  } yield GeneratedMasterNodeSeed(
    r.upvestMsgId,
    r.userId,
    r.userSecret,
    mns.seedHash,
    mns,
    assetKinds = r.desiredAssetKinds getOrElse SupportedAssetKinds
      intersect SupportedAssetKinds,
    r.callback,
  )

  def reencryptMasterNodeSeed(
    r: ReencryptMasterNodeSeed
  ): F[ReencryptedMasterNodeSeed] = for {
    mns0 <- loadMasterNodeSeedBlob(r)(r.seedHash)
    mns1 <- liftBreakout(r)(
      breakout.mns.reencrypt(
        mns0.userSecretEncrypted, mns0.seedHash,
        masterNodeSeedKey.pubkey,
        current = r.userSecret,
        next = (r.newUserSecret map insertDefaultUserSecretKey) getOrElse
          r.userSecret,
        r.upvestMsgId,
        extractKdfParamsOverrides(r),
      )
    )
    _ <- storeMasterNodeSeed(r)(mns1)
    _ <- flogger.info("Re-encrypted master-node-seed", mns1, r.upvestMsgId)
  } yield ReencryptedMasterNodeSeed(
    r.upvestMsgId,
    seedHash = mns1.seedHash,
    newUserSecret = r.newUserSecret,
    mns1,
    callback = r.callback
  )

  def recoverMasterNodeSeed(
    r: RecoverMasterNodeSeed
  ): F[RecoveredMasterNodeSeed] = for {
    mns <- liftBreakout(r)(
      breakout.mns.recover(
        r.sealedMasterNodeSeed.sealedBox,
        r.sealedMasterNodeSeed.seedHash,
        masterNodeSeedKey,
        r.newUserSecret,
        r.upvestMsgId,
        extractKdfParamsOverrides(r),
      )
    )
    _ <- storeMasterNodeSeed(r)(mns)
    _ <- flogger.info("Recovered master-node-seed", mns, r.upvestMsgId)
  } yield RecoveredMasterNodeSeed(
    upvestMsgId = r.upvestMsgId,
    seedHash = mns.seedHash,
    newUserSecret = r.newUserSecret,
    callback = r.callback
  )

  private def derivePath(
    abs: Option[HD.Path], rel: Option[HD.Path], root: HD.Path
  ): HD.Path = abs getOrElse rel.foldLeft(root) { _ / _ }

  private def checkExpectedAddress[A](r: Request)(
    expected: Option[A], got: A
  ): F[Unit] =
    expected match {
      case Some(e) if e != got =>
        M raiseError Failure.UnexpectedAddress(r,
          expected = e.toString,
          got = got.toString
          )
      case _ => M unit
    }

  def generateEthereumWallet(
    r: GenerateEthereumWallet
  ): F[GeneratedEthereumWallet] = for {
    emns <- loadMasterNodeSeedBlob(r)(r.seedHash)
    absPath = derivePath(
      abs = r.absolutePath,
      rel = r.relativePath,
      root = AssetKind.Ethereum.root
    )
    w <- liftBreakout(r)(breakout.eth.generate(
      r.userSecret, emns.userSecretEncrypted, absPath, r.upvestMsgId,
      extractKdfParamsOverrides(r),
    ))
    _ <- checkExpectedAddress(r)(r.expectedAddress, w.address)
    _ <- flogger.info(
      "Generated an Ethereum wallet",
      w.address, r.seedHash, r.upvestMsgId, absPath
    )
  } yield GeneratedEthereumWallet(
    r.upvestMsgId,
    r.seedHash,
    r.userId,
    r.callback,
    w.address,
    w.pk,
    absPath
  )

  def generateArweaveWallet(
    r: GenerateArweaveWallet
  ): F[GeneratedArweaveWallet] = for {
    emns <- loadMasterNodeSeedBlob(r)(r.seedHash)
    absPath = derivePath(
      abs = r.absolutePath,
      rel = r.relativePath,
      root = AssetKind.Arweave.root
    )
    w <- liftBreakout(r)(breakout.ar.generate(
      r.userSecret, emns.userSecretEncrypted, absPath, r.upvestMsgId,
      extractKdfParamsOverrides(r),
    ))
    _ <- checkExpectedAddress(r)(r.expectedAddress, w.address)
    _ <- flogger.info(
      "Generated an Arweave wallet",
      w.address, r.seedHash, r.upvestMsgId, absPath
    )
  } yield GeneratedArweaveWallet(
    upvestMsgId = r.upvestMsgId,
    seedHash = r.seedHash,
    userId = r.userId,
    callback = r.callback,
    address = w.address,
    crypto = w.pk,
    absolutePath = absPath,
  )

  def generateBitcoinWallet(
    r: GenerateBitcoinWallet
  ): F[GeneratedBitcoinWallet] = for {
    emns <- loadMasterNodeSeedBlob(r)(r.seedHash)
    absPath = derivePath(
      abs = r.absolutePath,
      rel = r.relativePath,
      root = AssetKind.Bitcoin.root
    )
    w <- liftBreakout(r)(
      breakout.btc.generate(bitcoinNetwork)(
        r.userSecret, emns.userSecretEncrypted, absPath, r.upvestMsgId,
        extractKdfParamsOverrides(r),
      )
    )
    _ <- checkExpectedAddress(r)(r.expectedAddress, w.address)
    _ <- flogger.info(
      "Generated a Bitcoin wallet",
      w.address, r.seedHash, r.upvestMsgId, absPath
    )
  } yield GeneratedBitcoinWallet(
    upvestMsgId = r.upvestMsgId,
    seedHash = r.seedHash,
    userId = r.userId,
    callback = r.callback,
    address = w.address,
    crypto = w.pk,
    absolutePath = absPath,
  )
}