ARC4 in Scala

ARC4.scala

/*
 * Copyright (c) 2017 Yuki Ono
 * Licensed under the MIT License.
 */

package func_arc4

import scala.annotation.tailrec

object ARC4 extends App {
    // https://github.com/bcgit/bc-java/blob/master/core/src/main/java/org/bouncycastle/crypto/engines/RC4Engine.java
    class CommonARC4() {
        private val STATE_LENGTH = 256
        private val state: Array[Byte] = new Array[Byte](STATE_LENGTH)
        private var x = 0
        private var y = 0
        
        private def next(): Byte = {
            x = (x + 1) % 0xff
            y = (state(x) + y) & 0xff
            
            val temp = state(x)
            state(x) = state(y)
            state(y) = temp
            
            state((state(x) + state(y)) & 0xff) 
        }
        
        def process(b: Byte): Byte = {
            (b ^ next()).asInstanceOf[Byte]
        }
        
        private def process(input: Array[Byte], output: Array[Byte]) {
            for (i <- (0 until input.length))
                output(i) = (input(i) ^ next()).asInstanceOf[Byte]
        }
        
        private def isNullOrEmpty(bytes: Array[Byte]): Boolean = {
            bytes == null || bytes.length == 0 
        }
        
        def process(input: Array[Byte]): Array[Byte] = {  
            require(!isNullOrEmpty(input))
            
            val output = new Array[Byte](input.length)
            process(input, output)
            
            output
        }
        
        def init(key: Array[Byte]) {
            require(!isNullOrEmpty(key))
            
            x = 0
            y = 0
            
            for (i <- (0 until STATE_LENGTH))
                state(i) = i.asInstanceOf[Byte]
            
            var i1 = 0
            var i2 = 0
            
            for (i <- (0 until STATE_LENGTH)) {
                i2 = ((key(i1) % 0xff) + state(i) + i2) & 0xff
                val temp = state(i)
                state(i) = state(i2)
                state(i2) = temp
                i1 = (i1 + 1) % key.length
            }
            
        }
    }
    
    case class ARC4Ctx(state: Array[Byte], x: Int, y: Int)
    
    object FuncARC4 {
        private val STATE_LENGTH = 256
        
        private def next(ctx: ARC4Ctx): (ARC4Ctx, Byte) = {
            val state = ctx.state
            val cx = (ctx.x + 1) % 0xff
            val cy = (state(cx) + ctx.y) & 0xff
            
            val temp = state(cx)
            state(cx) = state(cy)
            state(cy) = temp
            
            val res = state((state(cx) + state(cy)) & 0xff) 
            
            (ARC4Ctx(state, cx, cy), res) 
        }
        
        def process(ctx: ARC4Ctx, b: Byte): (ARC4Ctx, Byte) = {
            val (nextCtx, nextB) = next(ctx)
            (nextCtx, (b ^ nextB).asInstanceOf[Byte])
        }
        
        private def process(ctx: ARC4Ctx, input: Array[Byte], output: Array[Byte]): ARC4Ctx = {
            @tailrec
            def loop(ctx: ARC4Ctx, i: Int): ARC4Ctx = {
                if (i >= input.length)
                    ctx
                else {
                    val (nextCtx, b) = process(ctx, input(i))
                    output(i) = b
                    loop(nextCtx, i + 1)
                }
            }
            
            loop(ctx, 0)
        }
        
        private def isNullOrEmpty(bytes: Array[Byte]): Boolean = {
            bytes == null || bytes.length == 0 
        }
        
        def process(ctx: ARC4Ctx, input: Array[Byte]): (ARC4Ctx, Array[Byte]) = {
            require(!isNullOrEmpty(input))
            
            val output = new Array[Byte](input.length)
            val nextCtx = process(ctx, input, output)
            
            (nextCtx, output)
        }
        
        def init(key: Array[Byte]): ARC4Ctx = {
            require(!isNullOrEmpty(key))
            
            val state = (0 until STATE_LENGTH).map(_.toByte).toArray
            val keyStream = Stream.continually(key.toStream).flatten.take(STATE_LENGTH)
            
            var i2 = 0
            for((i, k) <- (0 until STATE_LENGTH).zip(keyStream)) {
                i2 = ((k & 0xff) + state(i) + i2) & 0xff
                val temp = state(i)
                state(i) = state(i2)
                state(i2) = temp
            }
            
            ARC4Ctx(state, 0, 0)
        }
    }
    
    val key = "key".getBytes
    //val key = "test".getBytes
 
    val carc4 = new CommonARC4()
    carc4.init(key)
    
    val emptyBytes = new Array[Byte](128)
    val cRes = carc4.process(emptyBytes)
    val cResHex = Hex.fpEncode(cRes)
    
    println("CommonARC4: " + cResHex)
  
    
    val ctx = FuncARC4.init(key)
    val emptyBytes2 = new Array[Byte](128)
    val (nextCtx, fRes) = FuncARC4.process(ctx, emptyBytes)
    val fResHex = Hex.fpEncode(fRes)
    
    println("FuncARC4: " + fResHex)
    
    println("\nOK: " + (cResHex == fResHex))
}