bitcoin_script_playground.swift

//: Bitcoin Script Interpreter

import Foundation

/** Script Operation
 */
enum Operation {
    case checkMultisig
    case checkSig
    case dup
    case equal
    case equalVerify
    case hash160
    case push(ScriptData)

    var asString: String {
        switch self {
        case .checkMultisig:
            return "OP_CHECKMULTISIG"
            
        case .checkSig:
            return "OP_CHECKSIG"
            
        case .dup:
            return "OP_DUP"
            
        case .equal:
            return "OP_EQUAL"
            
        case .equalVerify:
            return "OP_EQUALVERIFY"
            
        case .hash160:
            return "OP_HASH160"
            
        case .push(let scriptData):
            switch scriptData {
            case .number(let n):
                return "OP_" + String(n)
                
            case .operations(let operations):
                return "<" + stringify(operations: operations) + ">"
                
            case .raw(let str):
                return str
            }
        }
    }
}

/** Check for equality between script operations
 */
func ==(lhs: Operation, rhs: Operation) -> Bool {
    switch (lhs, rhs) {
    case (.checkMultisig, .checkMultisig),
         (.checkSig, .checkSig),
         (.dup, .dup),
         (.equal, .equal),
         (.equalVerify, .equalVerify),
         (.hash160, .hash160):
        return true
        
    case (.push(let a), .push(let b)):
        return a == b
        
    default:
        return false
    }
}

/** Script data pushes
*/
enum ScriptData {
    case number(Int)
    case operations([Operation])
    case raw(String)
}

/** Check for equality between script datas
 */
func ==(lhs: ScriptData, rhs: ScriptData) -> Bool {
    switch (lhs, rhs) {
    case (.number(let a), .number(let b)):
        return a == b
        
    case (.raw(let a), .raw(let b)):
        return a == b
        
    default:
        return false
    }
}

/** Script Evaluation Errors
 */
enum OperationError: Error {
    case checkSigFailed
    case equalityCheckFailed
    case equalVerifyFailed
    case expectedMoreStackItems
    case expectedMultisigPubKey
    case expectedMultisigSignature
    case expectedMultisigSignaturesCount
    case expectedMultisigTotalCount
    case expectedNullDummy
    case expectedSerializedScript
    case unexpectedHashNumberOperation
}

/** Get a series of operations as a familiar looking string
 */
func stringify(operations: [Operation]) -> String {
    return operations.map { $0.asString }.joined(separator: " ")
}

/** Determine if a sequence of operations matches the pay to script hash pattern
*/
func isPayToScriptHash(operations: [Operation]) -> Bool {
    let payToScriptHashPattern: [Operation] = [
        .hash160,
        .push(ScriptData.raw("scriptHash")),
        .equal
    ]
    
    guard operations.count >= payToScriptHashPattern.count else { return false }
    
    let presentPatternOperations = (0..<payToScriptHashPattern.count).filter { i in
        return operations[operations.count - i - 1] == payToScriptHashPattern[payToScriptHashPattern.count - i - 1]
    }
    
    return presentPatternOperations.count == payToScriptHashPattern.count
}

/** Evaluate a sequence of operations to confirm authorization
 */
func run(_ operations: [Operation]) throws {
    var stack = [ScriptData]()

    try operations.enumerated().forEach { index, op in
        let isFinalOperation = index == operations.count - 1

        switch op {
        case .checkMultisig:
            // Pop the total pub keys count from the stack
            guard
                let multisigTotalCountData = stack.popLast(),
                case .number(let multisigTotalCount) = multisigTotalCountData
                else
            {
                throw OperationError.expectedMultisigTotalCount
            }
            
            try (0..<multisigTotalCount).forEach { i in
                // Pop the pub keys from the stack
                guard let pubKey = stack.popLast() else { throw OperationError.expectedMoreStackItems }
                
                guard pubKey == ScriptData.raw("pubKey" + String(multisigTotalCount - i)) else {
                    throw OperationError.expectedMultisigPubKey
                }
            }
            
            // Pop the total signatures present from the stack
            guard
                let multisigSignaturesCountData = stack.popLast(),
                case .number(let signaturesCount) = multisigSignaturesCountData
                else
            {
                throw OperationError.expectedMultisigSignaturesCount
            }
            
            try (0..<signaturesCount).forEach { i in
                // Pop the signatures from the stack
                guard let sig = stack.popLast() else { throw OperationError.expectedMoreStackItems }
                
                guard sig == ScriptData.raw("sig" + String(signaturesCount - i)) else {
                    throw OperationError.expectedMultisigSignature
                }
            }
            
            guard let dummy = stack.popLast(), case .number(let dummyNumber) = dummy, dummyNumber == 0 else {
                throw OperationError.expectedNullDummy
            }
            
        case .checkSig:
            let items = (pubKey: stack.popLast(), sig: stack.popLast())
            
            guard isFinalOperation else { break }
            
            guard
                let pubKey = items.pubKey,
                let sig = items.sig,
                pubKey == ScriptData.raw("pubKey"),
                sig == ScriptData.raw("sig")
                else
            {
                throw OperationError.checkSigFailed
            }
            
        case .dup:
            if let lastStackItem = stack.last { stack += [lastStackItem] }
            
        case .equal:
            guard let ultimate = stack.popLast(), let penultimate = stack.popLast() else {
                throw OperationError.expectedMoreStackItems
            }
            
            guard isFinalOperation else { break }
            
            guard ultimate == penultimate else { throw OperationError.equalityCheckFailed }

        case .equalVerify:
            let items = (ultimate: stack.popLast(), penultimate: stack.popLast())
            
            guard
                let ultimate = items.ultimate,
                let penultimate = items.penultimate,
                ultimate == penultimate
                else
            {
                throw OperationError.equalVerifyFailed
            }

        case .hash160:
            guard let scriptData = stack.popLast() else { throw OperationError.expectedMoreStackItems }
            
            switch scriptData {
            case .number(_):
                throw OperationError.unexpectedHashNumberOperation
                
            case .operations(_):
                // Pretend to hash the script
                stack += [ScriptData.raw("scriptHash")]
                
            case .raw(_):
                // Pretend to hash the pubKey
                stack += [ScriptData.raw("pubKeyHash")]
            }
            
        case .push(let scriptData):
            stack += [scriptData]
        }
    }

    if isPayToScriptHash(operations: operations) {
        let signatures = Array(operations[0..<operations.count - 4])
        
        let nullDummy = Operation.push(ScriptData.number(0))
        
        let scriptOperation = operations[operations.count - 4]
        
        guard
            case .push(let scriptOperationsData) = scriptOperation,
            case .operations(let scriptOperations) = scriptOperationsData
            else
        {
            throw OperationError.expectedSerializedScript
        }
        
        let payToScriptHashSecondPassOperations = [nullDummy] + signatures + scriptOperations
        
        try run(payToScriptHashSecondPassOperations)
    }
}

protocol Script {
    var outputAsm: [Operation] { get }
    var redeemScript: [Operation] { get }
}

extension Script {
    /** End to end script
     */
    var script: [Operation] { return redeemScript + outputAsm }
}

struct PayToPublicKeyHash: Script {
    /** scriptPubKey / output
    */
    let outputAsm: [Operation] = [
        .dup,
        .hash160,
        .push(ScriptData.raw("pubKeyHash")),
        .equalVerify,
        .checkSig
    ]
    
    /** redeem script / input
    */
    let redeemScript: [Operation] = [
        .push(ScriptData.raw("sig")),
        .push(ScriptData.raw("pubKey"))
    ]
}

struct PayToScriptHash: Script {
    /** scriptPubKey / output
    */
    let outputAsm: [Operation] = [
        .hash160,
        .push(ScriptData.raw("scriptHash")),
        .equal
    ]
    
    /** redeemScript / input
    */
    let redeemScript: [Operation] = [
        .push(ScriptData.raw("sig1")),
        .push(ScriptData.raw("sig2")),
        .push(ScriptData.operations([
            .push(ScriptData.number(2)),
            .push(ScriptData.raw("pubKey1")),
            .push(ScriptData.raw("pubKey2")),
            .push(ScriptData.raw("pubKey3")),
            .push(ScriptData.number(3)),
            .checkMultisig
        ]))
    ]
}

do {
    let p2pkh = PayToPublicKeyHash().script
    
    print("Evaluating " + stringify(operations: p2pkh))
    
    try run(p2pkh)
    
    let p2sh = PayToScriptHash().script
    
    print("Evaluating " + stringify(operations: p2sh))
    
    try run(p2sh)
    
    print("Successfully evaluated all scripts")
}
catch let err {
    print("Execution failed", err)
}