MTProtoTest.scala

import scala.io.StdIn
import scala.language.{implicitConversions, postfixOps}

import scodec.bits.ByteVector
import Types.{ReqDHParams, ReqPQ, ResPQ}
import akka.actor.ActorSystem
import akka.stream.scaladsl.Tcp
import akka.stream.ActorMaterializer
import akka.util.ByteString
import scodec.{Codec, DecodeResult}
import scodec.Attempt.{Failure, Successful}

object Implicits {
  implicit def byteStringToByteVector(bs: ByteString): ByteVector = ByteVector(bs.toArray)
  implicit def byteVectorToByteString(bs: ByteVector): ByteString = ByteString.fromArrayUnsafe(bs.toArray)
}

import Implicits._

object TestCrypto {
  import javax.crypto.Cipher
  import java.security.KeyPairGenerator

  val RSAKeyPair = {
    val keyGen = KeyPairGenerator.getInstance("RSA")
    keyGen.generateKeyPair()
  }

  def decodeRSA(str: String): String = {
    val cipher = Cipher.getInstance("RSA/ECB/NoPadding")
    cipher.init(Cipher.DECRYPT_MODE, RSAKeyPair.getPrivate)
    new String(cipher.doFinal(str.getBytes), "ASCII")
  }
}

object Types {
  type Nonce = ByteVector

  import scodec._
  import codecs._
  import scodec.bits._

  final case class ReqPQ(nonce: Nonce)
  object ReqPQ {
    implicit val codec = (constant(hex"60469778") ~> bytes(16)).as[ReqPQ]
  }

  final case class ResPQ(nonce: Nonce, serverNonce: Nonce, pq: String, serverPublicKeyFingerprints: List[Long])
  object ResPQ {
    implicit val codec = (constant(hex"05162463") ~> bytes(16) :: bytes(16) :: ascii32 :: listOfN(int32, int64)).as[ResPQ]
  }

  final case class ReqDHParams(nonce: Nonce, serverNonce: Nonce, p: String, q: String, publicKeyFingerPrint: Long, encryptedData: String)
  object ReqDHParams {
    implicit val codec = (constant(hex"d712e4be") ~> bytes(16) :: bytes(16) :: ascii32 :: ascii32 :: int64 :: ascii32).as[ReqDHParams]
  }
}

object Stages {
  import akka.NotUsed
  import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
  import akka.stream.scaladsl.Flow
  import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
  import akka.util.ByteString

  object MTProtoTestStage {
    def apply(): Flow[ByteString, ByteString, NotUsed] = {
      Flow.fromGraph(new MTProtoTestStage())
    }
  }

  // MTProto TCP transport wrapper fields, CRC32 etc is not used
  class MTProtoTestStage() extends GraphStage[FlowShape[ByteString, ByteString]] {
    val inlet = Inlet[ByteString]("MTProtoTestStage.in")
    val outlet = Outlet[ByteString]("MTProtoTestStage.out")
    val shape = FlowShape(inlet, outlet)

    def createLogic(inheritedAttributes: Attributes) = new GraphStageLogic(shape) {
      class StageHandlers(inlet: Inlet[ByteString], outlet: Outlet[ByteString]) {
        var buffer = ByteString.empty

        class BufferedBytesHandler extends InHandler {
          override def onPush(): Unit = {
            val bytes = grab(inlet)
            buffer ++= bytes
            pull(inlet)
          }
        }

        def processBufferMessage[T: Codec](onMessage: T ⇒ Unit): Unit = {
          val request = implicitly[Codec[T]].decode(buffer.bits)
          request match {
            case Successful(DecodeResult(value, remainder)) ⇒
              onMessage(value)
              buffer = remainder.bytes

            case _ ⇒ // Pass
          }
        }
        
        final class WaitForReqPQ extends BufferedBytesHandler {
          override def onPush(): Unit = {
            super.onPush()
            processBufferMessage { t: ReqPQ ⇒
              PushOperations.sendResPQ(t)
              StageTransitions.waitReqDHParams(t)
            }
          }
        }

        final class WaitForReqDHParams extends BufferedBytesHandler {
          override def onPush(): Unit = {
            super.onPush()
            processBufferMessage { t: ReqDHParams ⇒
              val result = TestCrypto.decodeRSA(t.encryptedData)
              println(result)
              StageTransitions.endStage()
            }
          }
        }

        object PushOperations {
          def sendResPQ(req: ReqPQ) = {
            val response = ResPQ(req.nonce, req.nonce.reverse, "123456", List(123456L))
            ResPQ.codec.encode(response) match {
              case Successful(value) ⇒
                emit(outlet, value.bytes: ByteString)
              case Failure(cause) ⇒
                failStage(new Exception(cause.toString()))
            }
          }
        }

        object StageTransitions {
          def waitForReqPQ(): Unit = {
            setHandler(inlet, new WaitForReqPQ)
          }

          def waitReqDHParams(reqPQ: ReqPQ): Unit = {
            setHandler(inlet, new WaitForReqDHParams)
          }

          def endStage(): Unit = {
            complete(outlet)
            cancel(inlet)
          }
        }
      }

      val handlersLogic = new StageHandlers(inlet, outlet)

      setHandler(outlet, new OutHandler {
        def onPull(): Unit = tryPull(inlet)
      })

      handlersLogic.StageTransitions.waitForReqPQ()
    }
  }
}

object MTProtoTest extends App {
  implicit val actorSystem = ActorSystem("test")
  implicit val materializer = ActorMaterializer()
  Tcp().bindAndHandle(Stages.MTProtoTestStage(), "0.0.0.0", 1234)
  StdIn.readLine("Press enter to exit")
  actorSystem.terminate()
}