CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
# -*- coding: binary -*-
2

3
module Rex
4
  module Crypto
5
    class Chacha20
6
      attr_reader :chacha_ctx
7

8
      def initialize(key, iv)
9
        raise TypeError unless key.is_a? String
10
        raise TypeError unless iv.is_a? String
11

12
        @chacha_ctx = chacha20_ctx_setup(key, iv)
13
      end
14

15
      def reset_cipher(key, iv)
16
        @chacha_ctx = chacha20_ctx_setup(key, iv)
17
      end
18

19
      def chacha20_update_ctx
20
        @chacha_ctx[12] = (@chacha_ctx[12] + 1) & 0xffffffff
21
      end
22

23
      def chacha20_ctx_setup(key, iv, position=1)
24
        chacha_constants = [1634760805, 857760878, 2036477234, 1797285236]
25
        chacha_ctx = chacha_constants
26

27
        chacha_ctx += key.unpack('V8')
28
        chacha_ctx[12] = position
29
        chacha_ctx += iv.unpack('VVV')
30

31
        chacha_ctx
32
      end
33

34
      def chacha20_crypt(data)
35
        num_blocks = data.length / 64
36
        if data.length % 64 != 0
37
         num_blocks += 1
38
        end
39

40
        keystream = []
41
        (1..num_blocks).each do |block|
42
          keystream << chacha20_block_func
43
          chacha20_update_ctx
44
        end
45
        keystream = keystream.flatten
46

47
        enc = []
48
        i = 0
49
        data.unpack("c*").each do |a|
50
          enc << (a.ord ^ keystream[i])
51
          i += 1
52
        end
53
        enc.pack("c*").force_encoding('ASCII-8BIT')
54
      end
55

56
      def chacha20_block_func
57
        chacha_state = @chacha_ctx.dup
58

59
        (1..10).each do |round|
60
          quarter_round(chacha_state, 0, 4, 8, 12)
61
          quarter_round(chacha_state, 1, 5, 9, 13)
62
          quarter_round(chacha_state, 2, 6, 10, 14)
63
          quarter_round(chacha_state, 3, 7, 11, 15)
64
          quarter_round(chacha_state, 0, 5, 10, 15)
65
          quarter_round(chacha_state, 1, 6, 11, 12)
66
          quarter_round(chacha_state, 2, 7, 8, 13)
67
          quarter_round(chacha_state, 3, 4, 9, 14)
68
        end
69

70
        keystream_arr = []
71
        (0..15).each do |index|
72
          keystream_ch = (chacha_state[index] + @chacha_ctx[index]) & 0xffffffff
73
          keystream_arr << (keystream_ch & 0xff)
74
          keystream_arr << (keystream_ch >> 8 & 0xff)
75
          keystream_arr << (keystream_ch >> 16 & 0xff)
76
          keystream_arr << (keystream_ch >> 24 & 0xff)
77
        end
78

79
        keystream_arr
80
      end
81

82
      def rotate(v, c)
83
        ((v << c) & 0xffffffff) | v >> (32 - c)
84
      end
85

86
      def quarter_round(x, a, b, c, d)
87
        x[a] = (x[a] + x[b]) & 0xffffffff
88
        x[d] = rotate(x[d] ^ x[a], 16)
89
        x[c] = (x[c] + x[d]) & 0xffffffff
90
        x[b] = rotate(x[b] ^ x[c], 12)
91
        x[a] = (x[a] + x[b]) & 0xffffffff
92
        x[d] = rotate(x[d] ^ x[a], 8)
93
        x[c] = (x[c] + x[d]) & 0xffffffff
94
        x[b] = rotate(x[b] ^ x[c], 7)
95
      end
96
    end
97
  end
98
end
99

100