udp/aes.go

   1 package udpapi
   2
   3 import (
   4         "crypto/aes"
   5         "crypto/cipher"
   6         "crypto/md5"
   7 )
   8
   9 // Yes, AniDB works in ECB mode
  10 type ecbState struct {
  11         cipher.Block
  12 }
  13
  14 func newECBState(udpKey string, salt []byte) *ecbState {
  15         ecb := &ecbState{}
  16         ecb.Init(udpKey, salt)
  17         return ecb
  18 }
  19
  20 func (ecb *ecbState) Init(udpKey string, salt []byte) {
  21         h := md5.New()
  22         h.Write([]byte(udpKey))
  23         h.Write(salt)
  24
  25         key := h.Sum(nil)
  26
  27         b, err := aes.NewCipher(key)
  28         ecb.Block = b
  29         if err != nil {
  30                 panic(err)
  31         }
  32 }
  33
  34 func (ecb *ecbState) BlockSize() int {
  35         return aes.BlockSize
  36 }
  37
  38 func (ecb *ecbState) Encrypt(p []byte) (c []byte) {
  39         if ecb == nil {
  40                 return p
  41         }
  42
  43         padded := pkcs7Pad(p, aes.BlockSize)
  44         c = make([]byte, len(padded))
  45
  46         for i := 0; i < len(padded); i += aes.BlockSize {
  47                 ecb.Block.Encrypt(c[i:i+aes.BlockSize], padded[i:i+aes.BlockSize])
  48         }
  49         return c
  50 }
  51
  52 func (ecb *ecbState) Decrypt(c []byte) []byte {
  53         if ecb == nil {
  54                 return c
  55         }
  56
  57         p := make([]byte, len(c))
  58         for i := 0; i < len(c); i += aes.BlockSize {
  59                 ecb.Block.Decrypt(p[i:i+aes.BlockSize], c[i:i+aes.BlockSize])
  60         }
  61         return pkcs7Unpad(p, aes.BlockSize)
  62 }
  63
  64 // examples for a blocksize of 4
  65 // "almost1\x1"
  66 // "bytes\x3\x3\x3"
  67 // "byte\x4\x4\x4\x4"
  68 func pkcs7Pad(b []byte, blockSize byte) (padded []byte) {
  69         ps := int(blockSize) - len(b)%int(blockSize)
  70         padded = make([]byte, 0, len(b)+ps)
  71         padded = append(padded, b...)
  72
  73         for i := 0; i < ps; i++ {
  74                 padded = append(padded, byte(ps))
  75         }
  76         return padded
  77 }
  78
  79 func pkcs7Unpad(b []byte, blockSize byte) (unpadded []byte) {
  80         ps := b[len(b)-1]
  81         if ps > blockSize {
  82                 return b
  83         }
  84         padding := b[len(b)-int(ps):]
  85         for _, pb := range padding {
  86                 if pb != ps {
  87                         return b
  88                 }
  89         }
  90         return b[:len(b)-int(ps)]
  91 }