An implementation of ARC4 for Nim

arc4.nim

## Nim implementation of ARC4 stream cipher
## https://en.wikipedia.org/wiki/RC4
## http://www.users.zetnet.co.uk/hopwood/crypto/scan/cs.html#RC4

import future


type ARC4* = object
  S: seq[int]
  keystream: iterator(S: var seq[int]): int

proc KSA(key: string): seq[int] =
  ## RC4 key scheduling algorithm, returns a state array based on the key
  let keylength = key.len
  var
    j = 0
    S = lc[x | (x <- 0..255), int] # intialize to identity permutation
  for i in 0..255:
    j = (j + S[i] + ord(key[i %% keylength])) %% 256
    swap S[i], S[j]
  return S

iterator PRGA(S: var seq[int]): int {.closure.} =
  ## Pseudo-random generation algorithm. This iterator modifies the state and
  ## outputs a byte of the keystream, the yeild of this iterator must be xor'ed with
  ## a byte of the plaintext in order to produce a byte of ciphertext
  var i, j = 0
  while true:
    i = (i + 1) %% 256
    j = (j + S[i]) %% 256
    swap S[i], S[j]
    yield S[(S[i] + S[j]) %% 256]

proc newARC4*(key: string, drop: int = 0): ARC4 =
  ## Constructor for a new ARC4 object, specify a key and optionally a drop int
  ## for RC4-drop[(nbytes)] based on "(Not So) Random Shuffles of RC4" by
  ## Ilya Mironov. 768 bytes drop is recommended as a "as reasonable a precaution"
  ## see: http://eprint.iacr.org/2002/067
  var
    i = drop
  result = ARC4(S: KSA(key),  keystream: PRGA)
  while i > 0:
    discard result.keystream(result.S)
    dec i

proc crypt*(self: var ARC4, text: string): string =
  ## Performs encryption/decryption when passed an initialized ARC4 object.
  result = ""
  for c in text:
    result.add(chr(ord(c) xor self.keystream(self.S)))


# Example use

when isMainModule:
  import strutils
  var
    test_vec1 = "Attack at dawn"
    test_vec2 = "Whatever you do take care of your shoes"

  stdout.write "Test vector 1 plaintext: '" & test_vec1 & "'\n"
  stdout.write "Eencrypted (hex): "
  var foo = newARC4("Secret")
  # should produce the ciphertext: 45A01F645FC35B383552544B9BF5
  for c in foo.crypt(test_vec1):
    stdout.write c.toHex

  # example using RC4-drop[n], see newARC4 comments above
  stdout.write "\n\nTest vector 2 plaintext: '" & test_vec2 & "'\n"
  stdout.write "Encrypted (hex): "
  var
    bar = newARC4("Secret", 768) # drop 768 bytes
    tmp = bar.crypt(test_vec2)

  for c in tmp:
    stdout.write c.toHex
  stdout.write "\n"

  # decryption example
  stdout.write "Decryption result: "
  bar = newARC4("Secret", 768)
  for c in bar.crypt(tmp):
    stdout.write c
  stdout.write "\n\n"

Constructive criticism more than welcome, this was done for fun and learning only!