main.swift

import Foundation
import Command
import Crypto

struct Signature: Codable {
    let signature: Data
}

struct KeyPair: Codable {
    let priv: Data
    let pub: Data
    init() {
        var initial_priv: [UInt8] = []
        var initial_pub: [UInt8] = []
        for _ in 0...65 {
            var subkey: [UInt8] = []
            for _ in 0...63 {
                subkey.append(UInt8.random(in:0...255))
            }

            initial_priv.append(contentsOf: subkey)

            var hash = Data(bytes: subkey)
            for _ in 1...255 {
                hash = try! SHA512.hash(hash)
            }
            initial_pub.append(contentsOf: hash)
        }
        priv = Data(bytes: initial_priv)
        pub = Data(bytes: initial_pub)
    }
}

struct BChain {
    public static func generateKeyPair() -> KeyPair {
        return KeyPair()
    }

    public static func sign(_ hash: Data, withKeyPair keypair: KeyPair) -> Data {
        var signature: [UInt8] = []
        var distance = 255*64
        for i in 0...65 {
            let byte: UInt8
            switch i {
                case 64:
                    byte = numericCast(distance/256)
                case 65:
                    byte = numericCast(distance-distance/256*256)
                default:
                    byte = hash[i]
            }
            let subkey = keypair.priv[(i*64)...(i*64+63)]
            var subsig = subkey
            for _ in 0..<byte {
                subsig = try! SHA512.hash(subsig)
                if i < 64 {
                    distance -= 1
                }
            }
            signature.append(contentsOf: subsig)
        }
        return Data(bytes: signature)
    }

    public static func verify(_ sig: Signature, hash: Data, key keypair: KeyPair) -> Bool {
        var value: [UInt8] = []
        var distance: Int = 0
        for i in 0...65 {
            let subsig = sig.signature[(i*64)...(i*64+63)]
            let target = keypair.pub[(i*64)...(i*64+63)]
            var currentSig: Data = subsig
            for j in 0...255 {
                if currentSig == target {
                    value.append(numericCast(255 - j))
                    if (i < 64) {
                        distance += j
                    }
                    break;
                }
                do {
                    currentSig = try SHA512.hash(currentSig)
                } catch {
                    return false
                }
            }
            if i != value.count-1 {
                return false
            }
        }
        if distance != Int(value[64])*256+Int(value[65]) {
            return false
        }
        for i in 0...63 {
            if value[i] != hash[i] {
                return false
            }
        }
        return true
    }
}

struct BChainCommand: Command {
    var arguments: [CommandArgument] {
        return [
                .argument(name: "command", help: [
                                                    "keypair: generates a public and private keypair",
                                                    "sign: sign a file",
                                                    "verify: verify a file against a signature",
                                                 ]),
               ]
    }

    var options: [CommandOption] {
        return [
                .value(name: "output", short: "o", help: ["output file"]),
                .value(name: "key", short: "k", help: ["key file"]),
                .value(name: "input", short: "i", help: ["input file"]),
                .value(name: "signature", short: "s", help: ["signature file"]),
               ]
    }

    var help: [String] {
        return ["Does stuff.",
                "And more stuff?"
               ]
    }

    func run(using ctx: CommandContext) throws -> Future<Void> {
        let console = ctx.console
        switch ctx.arguments["command"] {

            case "keypair":
                guard let output = ctx.options["output"] else {
                    console.print("need an output file")
                    exit(-1)
                }
                let keypair = BChain.generateKeyPair()
                let encoder = JSONEncoder()
                encoder.outputFormatting = [JSONEncoder.OutputFormatting.prettyPrinted, .sortedKeys]
                let data = try! encoder.encode(keypair)
                let hash = try! SHA512.hash(keypair.pub)
                console.print("outputting keypair: \(hash.hexEncodedString())")
                let outputFilePath = URL(fileURLWithPath: output)
                do {
                    try data.write(to:outputFilePath)
                } catch {
                    console.print("Error writing key: \(error)")
                    exit(-1)
                }

            case "sign":
                //import key
                guard let key = ctx.options["key"] else {
                    console.print("need a key file")
                    exit(-1)
                }
                let decoder = JSONDecoder()
                let keypair: KeyPair
                do {
                    keypair = try decoder.decode(KeyPair.self, from: try Data(contentsOf: URL(fileURLWithPath: key)))
                } catch {
                    console.print("Error importing key: \(error)")
                    exit(-1)
                }

                //read input
                guard let input = ctx.options["input"] else {
                    console.print("need an input file")
                    exit(-1)
                }
                let data: Data
                do {
                    data = try Data(contentsOf: URL(fileURLWithPath: input))
                } catch {
                    console.print("Error loading input file: \(error)")
                    exit(-1)
                }

                //hash input
                let hash: Data
                do {
                    hash = try SHA512.hash(data)
                } catch {
                    console.print("Error hashing input: \(error)")
                    exit(-1)
                }

                //calculate signature
                let sig = Signature(signature: BChain.sign(hash, withKeyPair: keypair))

                //output signature
                guard let output = ctx.options["output"] else {
                    console.print("need an output file")
                    exit(-1)
                }
                let encoder = JSONEncoder()
                encoder.outputFormatting = [JSONEncoder.OutputFormatting.prettyPrinted, .sortedKeys]
                let json = try! encoder.encode(sig)
                let outputFilePath = URL(fileURLWithPath: output)
                do {
                    try json.write(to:outputFilePath)
                } catch {
                    console.print("Error writing key: \(error)")
                    exit(-1)
                }

            case "verify":
                //import key
                guard let key = ctx.options["key"] else {
                    console.print("need a key file")
                    exit(-1)
                }
                let decoder = JSONDecoder()
                let keypair: KeyPair
                do {
                    keypair = try decoder.decode(KeyPair.self, from: try Data(contentsOf: URL(fileURLWithPath: key)))
                } catch {
                    console.print("Error importing key: \(error)")
                    exit(-1)
                }

                //read input
                guard let input = ctx.options["input"] else {
                    console.print("need an input file")
                    exit(-1)
                }
                let data: Data
                do {
                    data = try Data(contentsOf: URL(fileURLWithPath: input))
                } catch {
                    console.print("Error loading input file: \(error)")
                    exit(-1)
                }

                //hash input
                let hash: Data
                do {
                    hash = try SHA512.hash(data)
                } catch {
                    console.print("Error hashing input: \(error)")
                    exit(-1)
                }

                //import signature
                guard let signature = ctx.options["signature"] else {
                    console.print("need a signature file")
                    exit(-1)
                }
                let sig: Signature
                do {
                    sig = try decoder.decode(Signature.self, from: try Data(contentsOf: URL(fileURLWithPath: signature)))
                } catch {
                    console.print("Error importing signature: \(error)")
                    exit(-1)
                }

                if BChain.verify(sig, hash: hash, key: keypair) {
                    console.print("verified")
                } else {
                    console.print("verification failed")
                    exit(-1)
                }


            default:
                console.print("Command is unrecognized.")
                exit(-1)
        }
        return .done(on: ctx.container)
    }
}

let console = Terminal()
let thread = MultiThreadedEventLoopGroup(numberOfThreads: 1)
var env = Environment(name: "Test", isRelease: false, arguments: CommandLine.arguments)
var container = BasicContainer(config: Config(), environment: env, services: Services(), on: thread)
do {
    try console.run(BChainCommand(), input: &env.commandInput, on: container).wait()
} catch {
    exit(-1)
}
exit(0)