DESAlgorithm.swift

//
//  main.swift
//  DESAlgorithmCommandLine
//
//  Created by Anton Poltoratskyi on 12.03.17.
//  Copyright © 2017 Anton Poltoratskyi. All rights reserved.
//

import Foundation

protocol St {}
protocol BitDecodable {
    func byteValue() -> UInt8
}
extension String: St {}
extension UInt8: BitDecodable {
    func byteValue() -> UInt8 {
        return self
    }
}

extension Int {
    func pow(times: Int) -> Int {
        var value = 1
        if times == 0 { return value }
        for _ in 1...times {
            value *= self
        }
        
        return value
    }
}

extension UInt8 {
    func pow(times: Int) -> UInt8 {
        var value: UInt8 = 1
        if times == 0 { return value }
        for _ in 1...times {
            value *= self
        }
        
        return value
    }
}

extension Array where Element: St {
    func toPlainInt() -> [Int] {
        return flatMap( { ($0 as! String)
            .components(separatedBy: " ") })
            .flatMap( {Int($0)} )
    }
}

extension Array where Element: BitDecodable {
    func shiftLeft(by times: Int) -> [Element] {
        var result = self
        for _ in 0..<times {
            result.append(result.removeFirst())
        }
        return result
    }
    
    func shiftRight(by times: Int) -> [Element] {
        var result = self
        for _ in 0..<times {
            result.insert(result.popLast()!, at: 0)
        }
        return result
    }
    
    func to64Blocks() -> [[UInt8]] {
        let countToAdd = count % 64 == 0 ? 0 : 64 - count % 64
        
        let addData = [UInt8](repeating: 0, count: countToAdd)
        let addicted = addData + self.map({ $0.byteValue()})
        let splitCount = addicted.count / 64
        
        var blocks: [[UInt8]] = []
        var currentBlock = 0
        
        var bytesArray = addicted.map { (byte) -> UInt8 in
            return byte
        }
        
        while currentBlock < splitCount {
            let subArray = bytesArray[currentBlock * 64..<(currentBlock + 1) * 64].map({$0})
            blocks.append(subArray)
            currentBlock += 1
        }
        
        return blocks
    }
    
    func toBytes() -> [UInt8] {
        assert(count % 8 == 0, "Wrong size")
        var bits = self
        var rawBytes: [[UInt8]] = []
        var byteIndex = 0
        
        while byteIndex < count / 8 {
            let byte = bits[byteIndex * 8...(byteIndex + 1) * 8 - 1].map { $0 as! UInt8 }
            rawBytes.append(byte)
            byteIndex += 1
        }
        
        return rawBytes.map { (bits) -> UInt8 in
            var value: UInt8 = 0
            
            bits.enumerated().forEach({ (index, bit) in
                value += bit * UInt8(2).pow(times: 7 - index)
            })
            
            return value
        }
    }
}

let IP = ["58 50 42 34 26 18 10 2 60 52 44 36 28 20 12 4",
          "62 54 46 38 30 22 14 6 64 56 48 40 32 24 16 8",
          "57 49 41 33 25 17 9 1 59 51 43 35 27 19 11 3",
          "61 53 45 37 29 21 13 5 63 55 47 39 31 23 15 7"].toPlainInt()

let decodeIP = ["40 8 48 16 56 24 64 32 39 7 47 15 55 23 63 31",
                "38 6 46 14 54 22 62 30 37 5 45 13 53 21 61 29",
                "36 4 44 12 52 20 60 28 35 3 43 11 51 19 59 27",
                "34 2 42 10 50 18 58 26 33 1 41 9 49 17 57 25"].toPlainInt()

let E = [
    "32 1 2 3 4 5",
    "4 5 6 7 8 9",
    "8 9 10 11 12 13",
    "12 13 14 15 16 17",
    "16 17 18 19 20 21",
    "20 21 22 23 24 25",
    "24 25 26 27 28 29",
    "28 29 30 31 32 1",
    ].toPlainInt()

let arrayS = [["14 4 13 1 2 15 11 8 3 10 6 12 5 9 0 7",
               "0 15 7 4 14 2 13 1 10 6 12 11 9 5 3 8",
               "4 1 14 8 13 6 2 11 15 12 9 7 3 10 5 0",
               "15 12 8 2 4 9 1 7 5 11 3 14 10 0 6 13"],
              [
                "15 1 8 14 6 11 3 4 9 7 2 13 12 0 5 10",
                "3 13 4 7 15 2 8 14 12 0 1 10 6 9 11 5",
                "0 14 7 11 10 4 13 1 5 8 12 6 9 3 2 15",
                "13 8 10 1 3 15 4 2 11 6 7 12 0 5 14 9"],
              
              ["10 0 9 14 6 3 15 5 1 13 12 7 11 4 2 8",
               "13 7 0 9 3 4 6 10 2 8 5 14 12 11 15 1",
               "13 6 4 9 8 15 3 0 11 1 2 12 5 10 14 7",
               "1 10 13 0 6 9 8 7 4 15 14 3 11 5 2 12"],
              
              ["7 13 14 3 0 6 9 10 1 2 8 5 11 12 4 15",
               "13 8 11 5 6 15 0 3 4 7 2 12 1 10 14 9",
               "10 6 9 0 12 11 7 13 15 1 3 14 5 2 8 4",
               "3 15 0 6 10 1 13 8 9 4 5 11 12 7 2 14"],
              
              ["2 12 4 1 7 10 11 6 8 5 3 15 13 0 14 9",
               "14 11 2 12 4 7 13 1 5 0 15 10 3 9 8 6",
               "4 2 1 11 10 13 7 8 15 9 12 5 6 3 0 14",
               "11 8 12 7 1 14 2 13 6 15 0 9 10 4 5 3"],
              
              ["12 1 10 15 9 2 6 8 0 13 3 4 14 7 5 11",
               "10 15 4 2 7 12 9 5 6 1 13 14 0 11 3 8",
               "9 14 15 5 2 8 12 3 7 0 4 10 1 13 11 6",
               "4 3 2 12 9 5 15 10 11 14 1 7 6 0 8 13"],
              
              ["4 11 2 14 15 0 8 13 3 12 9 7 5 10 6 1",
               "13 0 11 7 4 9 1 10 14 3 5 12 2 15 8 6",
               "1 4 11 13 12 3 7 14 10 15 6 8 0 5 9 2",
               "6 11 13 8 1 4 10 7 9 5 0 15 14 2 3 12"],
              
              ["13 2 8 4 6 15 11 1 10 9 3 14 5 0 12 7",
               "1 15 13 8 10 3 7 4 12 5 6 11 0 14 9 2",
               "7 11 4 1 9 12 14 2 0 6 10 13 15 3 5 8",
               "2 1 14 7 4 10 8 13 15 12 9 0 3 5 6 11"]
    ].map({$0.map({$0.components(separatedBy: " ")}).map({$0.flatMap({Int($0)})})})

//let S1 = arrayS[0]
//let S2 = arrayS[1]
//let S3 = arrayS[2]
//let S4 = arrayS[3]
//let S5 = arrayS[4]
//let S6 = arrayS[5]
//let S7 = arrayS[6]
//let S8 = arrayS[7]


let P = [
    "16 7 20 21 29 12 28 17",
    "1 15 23 26 5 18 31 10",
    "2 8 24 14 32 27 3 9",
    "19 13 30 6 22 11 4 25"
    ].toPlainInt()


protocol DataDecodable {
    func data() -> Data
}

extension String: DataDecodable {
    func data() -> Data {
        return data(using: .utf8)!
    }
}


extension Int {
    func bits() -> [UInt8] {
        let bits = [self & 0b1,
                    self >> 1 & 0b1,
                    self >> 2 & 0b1,
                    self >> 3 & 0b1]
        return bits.map { UInt8($0) }
    }
}

extension Data {
    func bits() -> [UInt8] {
        return self.flatMap { (byte) -> [UInt8] in
            let bits = [ Int(byte) >> 0 & 0b1,
                         Int(byte) >> 1 & 0b1,
                         Int(byte) >> 2 & 0b1,
                         Int(byte) >> 3 & 0b1,
                         Int(byte) >> 4 & 0b1,
                         Int(byte) >> 5 & 0b1,
                         Int(byte) >> 6 & 0b1,
                         Int(byte) >> 7 & 0b1]
            
            return bits.map({ UInt8($0) })
        }
    }
}

func replaceWithIP(ip: [Int], bytes: [UInt8]) -> [UInt8] {
    var swapped = [UInt8]()
    
    for (swapIndex) in ip {
        
        //print("Replace: ", bytes[swapIndex - 1])
        swapped.append(bytes[swapIndex - 1])
    }
    
    return swapped
}

infix operator ^^
public func ^^(lhs: [UInt8], rhs: [UInt8]) -> [UInt8] {
    assert(lhs.count == rhs.count)
    
    return zip(lhs, rhs).map { (r, k) -> UInt8 in
        return r ^ k
    }
}

//MARK: - Key generating
func extendWithBits(key: [UInt8]) -> [UInt8] {
    var batch = 0
    var result: [UInt8] = []
    
    while batch < 8 {
        var bytes = key[batch * 7...(batch + 1) * 7 - 1]
        let onesCount = bytes.filter({ $0 == 1 }).count
        bytes.append(onesCount % 2 == 0 ? 1 : 0)
        result.append(contentsOf: bytes)
        batch += 1
    }
    
    return result
}

let C0table = ["57 49 41 33 25 17 9 1 58 50 42 34 26 18",
               "10 2 59 51 43 35 27 19 11 3 60 52 44 36"].toPlainInt()
let D0table = ["63 55 47 39 31 23 15 7 62 54 46 38 30 22",
               "14 6 61 53 45 37 29 21 13 5 28 20 12 4"].toPlainInt()


func generateKeys(from primary: [UInt8], encode: Bool) -> [[UInt8]] {
    var keys: [[UInt8]] = []
    let shiftsTable = encode
        ? [1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 1]
        : [0, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 2, 1]
    
    let validPositions = [14, 17, 11, 24, 1, 5, 3, 28, 15, 6, 21, 10, 23, 19, 12, 4,
                          26, 8, 16, 7, 27, 20, 13, 2, 41, 52, 31, 37, 47, 55, 30, 40,
                          51, 45, 33, 48, 44, 49, 39, 56, 34, 53, 46, 42, 50, 36, 29, 32]
    var C0: [UInt8] = []
    var D0: [UInt8] = []
    
    zip(C0table, D0table).forEach { (cIndex, dIndex) in
        C0.append(primary[cIndex - 1])
        D0.append(primary[dIndex - 1])
    }
    var Ci: [[UInt8]] = []
    var Di: [[UInt8]] = []
    
    Ci.append(encode
        ? C0.shiftLeft(by: shiftsTable[0])
        : C0.shiftRight(by: shiftsTable[0]))
    
    Di.append(encode
        ? D0.shiftLeft(by: shiftsTable[0])
        : D0.shiftRight(by: shiftsTable[0]))
    
    for round in 2...16 {
        
        Ci.append(encode
            ? Ci[round - 2].shiftLeft(by: shiftsTable[round - 1])
            : Ci[round - 2].shiftRight(by: shiftsTable[round - 1]))
        
        Di.append(encode
            ? Di[round - 2].shiftLeft(by: shiftsTable[round - 1])
            : Di[round - 2].shiftRight(by: shiftsTable[round - 1]))
    }
    
    keys = zip(Ci, Di).map({ (c, d) -> [UInt8] in
        let fullKey = c + d
        var validKey: [UInt8] = []
        
        for (swapIndex) in validPositions {
            validKey.append(fullKey[swapIndex - 1])
        }
        return validKey
    })
    
    return keys
}

func feystel(R: [UInt8], key: [UInt8]) -> [UInt8] {
    
    var extended: [UInt8] = []
    for (swapIndex) in E {
        extended.append(R[swapIndex - 1])
    }
    
    var xored: [UInt8] = extended ^^ key
    var blocks: [[UInt8]] = []
    
    var blockIndex = 0
    
    while blockIndex < 8 {
        let block = xored[blockIndex * 6...(blockIndex + 1) * 6 - 1].map {$0}
        blocks.append(block)
        blockIndex += 1
    }
    
    func getA(block: [UInt8]) -> Int {
        return Int(block[0]).pow(times: 1) + Int(block[5]).pow(times: 0)
    }
    
    func getB(block: [UInt8]) -> Int {
        return
            Int(block[1]).pow(times: 3) +
                Int(block[2]).pow(times: 2) +
                Int(block[3]).pow(times: 1) +
                Int(block[4]).pow(times: 0)
    }
    
    var bitBlocks4: [[UInt8]] = []
    
    blocks.enumerated().forEach { (index, block) in
        let S = arrayS[index]
        let a = getA(block: block)
        let b = getB(block: block)
        let intValue = S[a][b]
        bitBlocks4.append(intValue.bits())
    }
    
    let block32bit = bitBlocks4.reduce([], +)
    
    var Pblock: [UInt8] = []
    for (swapIndex) in P {
        Pblock.append(block32bit[swapIndex - 1])
    }
    
    return Pblock
}



func encryptDES(data: DataDecodable, key: String) -> [UInt8] {
    
    let blocks = data.data().bits().to64Blocks()
    let ipReplacedData = blocks.map({ replaceWithIP(ip: IP, bytes: $0) })
    
    let key = key.data(using: .utf8)!.bits()
    assert(key.count == 56, "Wrong key")
    let extendedKey = extendWithBits(key: key)
//    let extendedKey: [UInt8] = [1, 1, 0, 1, 1, 1, 1, 0,
//                       0, 0, 0, 1, 0, 0, 0, 0,
//                       1, 0, 0, 1, 1, 1, 0, 0,
//                       0, 1, 0, 1, 1, 0, 0, 0,
//                       1, 1, 1, 0, 1, 0, 0, 0,
//                       1, 0, 1, 0, 0, 1, 0, 0, 
//                       1, 0, 1, 0, 0, 1, 1, 0, 
//                       0, 0, 1, 1, 0, 0, 0, 0]
    let roundKeys = generateKeys(from: extendedKey, encode: true)
    
    //MARK: - Algo
    return ipReplacedData
        .flatMap { (block) -> [UInt8] in
            
            let L0 = block[0..<32].map { $0 }
            let R0 = block[32..<64].map { $0 }
            
            var Li = L0
            var Ri = R0
            
            for round in 1...16 {
                let LiPrev = Li
                Li = Ri
                Ri = LiPrev ^^ feystel(R: Ri, key: roundKeys[round - 1])
            }
            
            
            return Li + Ri
        }.reduce([], +)
}



func decryptDES(_ encryptedData: [UInt8], key: String) -> [UInt8] {
    let blocks = encryptedData.to64Blocks()
    let ipReplacedData = blocks.map({ replaceWithIP(ip: decodeIP, bytes: $0) })
    
    let key = key.data(using: .utf8)!.bits()
    assert(key.count == 56, "Wrong key")
    let extendedKey = extendWithBits(key: key)
    let roundKeys = generateKeys(from: extendedKey, encode: false).reversed().map { $0 }
    
    
    //MARK: - Algo
    return ipReplacedData
        .flatMap { (block) -> [UInt8] in
            
            let L0 = block[0..<32].map { $0 }
            let R0 = block[32..<64].map { $0 }
            
            var Li = L0
            var Ri = R0
            
            for round in 1...16 {
                let RiPrev = Ri
                Ri = Li
                Li = RiPrev ^^ feystel(R: Li, key: roundKeys[round - 1])
            }
            
            
            return Li + Ri
        }.reduce([], +)
}

let encryptedData = encryptDES(data: "Anton", key: "Hello!!")
let encryptedBytes = encryptedData.toBytes()
let encryptedString = String(data: Data(bytes: encryptedBytes), encoding: .utf8)

let decryptedData = decryptDES(encryptedData, key: "Hello!!")
let bytes = decryptedData.toBytes()
let data = Data(bytes: bytes)
let string = String(data: data, encoding: .utf8)

let a = 1


let first: [UInt8] = [1,0,1,0]
let second: [UInt8] = [0,1,1,0]