import Foundation

extension Data {
  init?(base64RawEncoded: String) {
    if base64RawEncoded.hasSuffix("=") {
      return nil
    }
    var str = base64RawEncoded
    switch base64RawEncoded.count % 4 {
    case 2:
      str += "=="
    case 3:
      str += "="
    default:
      ()
    }
    guard let data = Data(base64Encoded: str) else {
      return nil
    }
    self = data
  }

  var base64RawEncodedData: Data {
    var s = base64EncodedData(options: [
      Base64EncodingOptions.lineLength64Characters, Base64EncodingOptions.endLineWithLineFeed,
    ])
    if let pi = s.firstIndex(of: Character("=").asciiValue!) {
      s = Data(s[s.startIndex..<pi])
    }
    return s
  }

  var base64RawEncodedString: String {
    return String(decoding: base64RawEncodedData, as: UTF8.self)
  }
}

// Spec: https://github.com/bitcoin/bips/blob/master/bip-0173.mediawiki
// Modified version of https://github.com/0xDEADP00L/Bech32

// Copyright 2018 Evolution Group Limited

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:

// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

import Foundation

/// Bech32 checksum implementation
public class Bech32 {
  private let gen: [UInt32] = [0x3b6a_57b2, 0x2650_8e6d, 0x1ea1_19fa, 0x3d42_33dd, 0x2a14_62b3]
  /// Bech32 checksum delimiter
  private let checksumMarker: String = "1"
  /// Bech32 character set for encoding
  private let encCharset: Data = Data("qpzry9x8gf2tvdw0s3jn54khce6mua7l".utf8)
  /// Bech32 character set for decoding
  private let decCharset: [Int8] = [
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
    15, -1, 10, 17, 21, 20, 26, 30, 7, 5, -1, -1, -1, -1, -1, -1,
    -1, 29, -1, 24, 13, 25, 9, 8, 23, -1, 18, 22, 31, 27, 19, -1,
    1, 0, 3, 16, 11, 28, 12, 14, 6, 4, 2, -1, -1, -1, -1, -1,
    -1, 29, -1, 24, 13, 25, 9, 8, 23, -1, 18, 22, 31, 27, 19, -1,
    1, 0, 3, 16, 11, 28, 12, 14, 6, 4, 2, -1, -1, -1, -1, -1,
  ]

  private func convertBits(from: Int, to: Int, pad: Bool, idata: Data) throws -> Data {
    var acc: Int = 0
    var bits: Int = 0
    let maxv: Int = (1 << to) - 1
    let maxAcc: Int = (1 << (from + to - 1)) - 1
    var odata = Data()
    for ibyte in idata {
      acc = ((acc << from) | Int(ibyte)) & maxAcc
      bits += from
      while bits >= to {
        bits -= to
        odata.append(UInt8((acc >> bits) & maxv))
      }
    }
    if pad {
      if bits != 0 {
        odata.append(UInt8((acc << (to - bits)) & maxv))
      }
    } else if bits >= from || ((acc << (to - bits)) & maxv) != 0 {
      throw DecodingError.bitsConversionFailed
    }
    return odata
  }

  /// Find the polynomial with value coefficients mod the generator as 30-bit.
  private func polymod(_ values: Data) -> UInt32 {
    var chk: UInt32 = 1
    for v in values {
      let top = (chk >> 25)
      chk = (chk & 0x1ffffff) << 5 ^ UInt32(v)
      for i: UInt8 in 0..<5 {
        chk ^= ((top >> i) & 1) == 0 ? 0 : gen[Int(i)]
      }
    }
    return chk
  }

  /// Expand a HRP for use in checksum computation.
  private func expandHrp(_ hrp: String) -> Data {
    let hrpBytes = Data(hrp.utf8)
    var result = Data(repeating: 0x00, count: hrpBytes.count * 2 + 1)
    for (i, c) in hrpBytes.enumerated() {
      result[i] = c >> 5
      result[i + hrpBytes.count + 1] = c & 0x1f
    }
    result[hrp.count] = 0
    return result
  }

  /// Verify checksum
  private func verifyChecksum(hrp: String, checksum: Data) -> Bool {
    var data = expandHrp(hrp)
    data.append(checksum)
    return polymod(data) == 1
  }

  /// Create checksum
  private func createChecksum(hrp: String, values: Data) -> Data {
    var enc = expandHrp(hrp)
    enc.append(values)
    enc.append(Data(repeating: 0x00, count: 6))
    let mod: UInt32 = polymod(enc) ^ 1
    var ret: Data = Data(repeating: 0x00, count: 6)
    for i in 0..<6 {
      ret[i] = UInt8((mod >> (5 * (5 - i))) & 31)
    }
    return ret
  }

  /// Encode Bech32 string
  private func encodeBech32(_ hrp: String, values: Data) -> String {
    let checksum = createChecksum(hrp: hrp, values: values)
    var combined = values
    combined.append(checksum)

    let hrpBytes = Data(hrp.utf8)
    var ret = hrpBytes
    ret.append(Data("1".utf8))
    for i in combined {
      ret.append(encCharset[Int(i)])
    }
    return String(decoding: ret, as: UTF8.self)
  }

  /// Decode Bech32 string
  public func decodeBech32(_ str: String) throws -> (hrp: String, checksum: Data) {
    let strBytes = Data(str.utf8)
    var lower: Bool = false
    var upper: Bool = false
    for c in strBytes {
      // printable range
      if c < 33 || c > 126 {
        throw DecodingError.nonPrintableCharacter
      }
      // 'a' to 'z'
      if c >= 97 && c <= 122 {
        lower = true
      }
      // 'A' to 'Z'
      if c >= 65 && c <= 90 {
        upper = true
      }
    }
    if lower && upper {
      throw DecodingError.invalidCase
    }
    guard let pos = str.range(of: checksumMarker, options: .backwards)?.lowerBound else {
      throw DecodingError.noChecksumMarker
    }
    let intPos: Int = str.distance(from: str.startIndex, to: pos)
    guard intPos >= 1 else {
      throw DecodingError.incorrectHrpSize
    }
    guard intPos + 7 <= str.count else {
      throw DecodingError.incorrectChecksumSize
    }
    let vSize: Int = str.count - 1 - intPos
    var values: Data = Data(repeating: 0x00, count: vSize)
    for i in 0..<vSize {
      let c = strBytes[i + intPos + 1]
      let decInt = decCharset[Int(c)]
      if decInt == -1 {
        throw DecodingError.invalidCharacter
      }
      values[i] = UInt8(decInt)
    }
    let hrp = String(str[..<pos]).lowercased()
    guard verifyChecksum(hrp: hrp, checksum: values) else {
      throw DecodingError.checksumMismatch
    }
    return (hrp, Data(values[..<(vSize - 6)]))
  }

  public func encode(hrp: String, data: Data) -> String {
    let isUpper = hrp[hrp.startIndex].isUppercase
    let result = encodeBech32(
      isUpper ? hrp.lowercased() : hrp,
      values: try! self.convertBits(from: 8, to: 5, pad: true, idata: data))
    return isUpper ? result.uppercased() : result
  }

  public func decode(_ str: String) throws -> (hrp: String, data: Data) {
    let isUpper = str[str.startIndex].isUppercase
    let result = try decodeBech32(isUpper ? str.lowercased() : str)
    return (
      isUpper ? result.hrp.uppercased() : result.hrp,
      try convertBits(from: 5, to: 8, pad: false, idata: result.checksum)
    )
  }
}

extension Bech32 {
  public enum DecodingError: LocalizedError {
    case nonUTF8String
    case nonPrintableCharacter
    case invalidCase
    case noChecksumMarker
    case incorrectHrpSize
    case incorrectChecksumSize

    case invalidCharacter
    case checksumMismatch

    case bitsConversionFailed

    public var errorDescription: String? {
      switch self {
      case .bitsConversionFailed:
        return "Failed to perform bits conversion"
      case .checksumMismatch:
        return "Checksum doesn't match"
      case .incorrectChecksumSize:
        return "Checksum size too low"
      case .incorrectHrpSize:
        return "Human-readable-part is too small or empty"
      case .invalidCase:
        return "String contains mixed case characters"
      case .invalidCharacter:
        return "Invalid character met on decoding"
      case .noChecksumMarker:
        return "Checksum delimiter not found"
      case .nonPrintableCharacter:
        return "Non printable character in input string"
      case .nonUTF8String:
        return "String cannot be decoded by utf8 decoder"
      }
    }
  }
}

import Foundation

let version = "v0.1.4"

@main
struct CLI {
  static func main() {
    do {
      let plugin = Plugin(crypto: CryptoKitCrypto(), stream: StandardIOStream())
      let options = try Options.parse(CommandLine.arguments)
      switch options.command {
      case .help:
        print(Options.help)
      case .version:
        print(version)
      case .keygen:
        let result = try plugin.generateKey(
          accessControl: options.accessControl.keyAccessControl,
          recipientType: options.recipientType.recipientType, now: Date())
        if let outputFile = options.output {
          FileManager.default.createFile(
            atPath: outputFile,
            contents: Data(result.0.utf8),
            attributes: [.posixPermissions: 0o600]
          )
          print("Public key: \(result.1)")
        } else {
          print(result.0)
        }
      case .recipients:
        var input = ""
        if let inputFile = options.input {
          input = try String(
            contentsOfFile: inputFile)
        } else {
          input = try String(data: FileHandle.standardInput.readToEnd()!, encoding: .utf8)!
        }
        let result = try plugin.generateRecipients(
          input: input, recipientType: options.recipientType.recipientType)
        if let outputFile = options.output {
          FileManager.default.createFile(
            atPath: outputFile,
            contents: Data(result.utf8),
            attributes: [.posixPermissions: 0o600]
          )
        } else if result != "" {
          print(result)
        }
      case .plugin(let sm):
        switch sm {
        case .recipientV1:
          plugin.runRecipientV1()
        case .identityV1:
          plugin.runIdentityV1()
        }
      }
    } catch {
      print("\(CommandLine.arguments[0]): error: \(error.localizedDescription)")
      exit(-1)
    }
  }
}

/// Command-line options parser
struct Options {
  enum Error: LocalizedError, Equatable {
    case unknownOption(String)
    case missingValue(String)
    case invalidValue(String, String)

    public var errorDescription: String? {
      switch self {
      case .unknownOption(let option): return "unknown option: `\(option)`"
      case .missingValue(let option): return "missing value for option `\(option)`"
      case .invalidValue(let option, let value):
        return "invalid value for option `\(option)`: `\(value)`"
      }
    }
  }

  enum StateMachine: String {
    case recipientV1 = "recipient-v1"
    case identityV1 = "identity-v1"
  }

  enum Command: Equatable {
    case help
    case version
    case keygen
    case recipients
    case plugin(StateMachine)
  }
  var command: Command

  var output: String?
  var input: String?

  enum AccessControl: String {
    case none = "none"
    case passcode = "passcode"
    case anyBiometry = "any-biometry"
    case anyBiometryOrPasscode = "any-biometry-or-passcode"
    case anyBiometryAndPasscode = "any-biometry-and-passcode"
    case currentBiometry = "current-biometry"
    case currentBiometryAndPasscode = "current-biometry-and-passcode"

    var keyAccessControl: KeyAccessControl {
      switch self {
      case .none: return KeyAccessControl.none
      case .passcode: return KeyAccessControl.passcode
      case .anyBiometry: return KeyAccessControl.anyBiometry
      case .anyBiometryOrPasscode: return KeyAccessControl.anyBiometryOrPasscode
      case .anyBiometryAndPasscode: return KeyAccessControl.anyBiometryAndPasscode
      case .currentBiometry: return KeyAccessControl.currentBiometry
      case .currentBiometryAndPasscode: return KeyAccessControl.currentBiometryAndPasscode
      }
    }
  }
  var accessControl = AccessControl.anyBiometryOrPasscode

  enum RecipientType: String {
    case se = "se"
    case p256tag = "p256tag"

    var recipientType: age_plugin_se.RecipientType {
      switch self {
      case .se: return .se
      case .p256tag: return .p256tag
      }
    }
  }

  var recipientType: RecipientType = .se

  static let help =
    """
    Usage:
      age-plugin-se keygen [-o OUTPUT] [--access-control ACCESS_CONTROL]
      age-plugin-se recipients [-o OUTPUT] [-i INPUT]

    Description:
      The `keygen` subcommand generates a new private key bound to the current 
      Secure Enclave, with the given access controls, and outputs it to OUTPUT 
      or standard output.

      The `recipients` subcommand reads an identity file from INPUT or standard 
      input, and outputs the corresponding recipient(s) to OUTPUT or to standard 
      output.

    Options:
      --access-control ACCESS_CONTROL   Access control for using the generated key.
                                        
                                        When using current biometry, adding or removing a 
                                        fingerprint stops the key from working. Removing an 
                                        added fingerprint enables the key again. 

                                        Supported values: none, passcode, 
                                          any-biometry, any-biometry-and-passcode, 
                                          any-biometry-or-passcode, current-biometry, 
                                          current-biometry-and-passcode
                                        Default: any-biometry-or-passcode.                          

      -i, --input INPUT                 Read data from the file at path INPUT

      -o, --output OUTPUT               Write the result to the file at path OUTPUT

    Example:
      $ age-plugin-se keygen -o key.txt
      Public key: age1se1qg8vwwqhztnh3vpt2nf2xwn7famktxlmp0nmkfltp8lkvzp8nafkqleh258
      $ tar cvz ~/data | age -r age1se1qgg72x2qfk9wg3wh0qg9u0v7l5dkq4jx69fv80p6wdus3ftg6flwg5dz2dp > data.tar.gz.age
      $ age --decrypt -i key.txt data.tar.gz.age > data.tar.gz
    """

  static func parse(_ args: [String]) throws -> Options {
    var opts = Options(command: .help)
    var i = 1
    while i < args.count {
      let arg = args[i]
      if arg == "keygen" {
        opts.command = .keygen
      } else if arg == "recipients" {
        opts.command = .recipients
      } else if ["--help", "-h"].contains(arg) {
        opts.command = .help
        break
      } else if ["--version"].contains(arg) {
        opts.command = .version
        break
      } else if [
        "--age-plugin", "-i", "--input", "-o", "--output", "--access-control", "--recipient-type",
      ].contains(where: {
        arg == $0 || arg.hasPrefix($0 + "=")
      }) {
        let argps = arg.split(separator: "=", maxSplits: 1)
        let value: String
        if argps.count == 1 {
          i += 1
          if i >= args.count {
            throw Error.missingValue(arg)
          }
          value = args[i]
        } else {
          value = String(argps[1])
        }
        let arg = String(argps[0])
        switch arg {
        case "--age-plugin":
          opts.command = try .plugin(
            StateMachine(rawValue: value) ?? { throw Error.invalidValue(arg, value) }())
        case "-i", "--input":
          opts.input = value
        case "-o", "--output":
          opts.output = value
        case "--access-control":
          opts.accessControl =
            try AccessControl(rawValue: value) ?? { throw Error.invalidValue(arg, value) }()
        case "--recipient-type":
          opts.recipientType =
            try RecipientType(rawValue: value) ?? { throw Error.invalidValue(arg, value) }()
        default:
          assert(false)
        }
      } else {
        throw Error.unknownOption(arg)
      }
      i += 1
    }
    return opts
  }
}

import Foundation

#if !os(Linux) && !os(Windows)
  import CryptoKit
  import LocalAuthentication
#else
  import Crypto
  struct SecAccessControl {}
#endif

/// Abstraction for random/unpredictable/system-specific crypto operations
protocol Crypto {
  var isSecureEnclaveAvailable: Bool { get }

  func newSecureEnclavePrivateKey(dataRepresentation: Data) throws -> SecureEnclavePrivateKey
  func newSecureEnclavePrivateKey(accessControl: SecAccessControl) throws -> SecureEnclavePrivateKey
  func newEphemeralPrivateKey() -> P256.KeyAgreement.PrivateKey
}

protocol SecureEnclavePrivateKey {
  var publicKey: P256.KeyAgreement.PublicKey { get }
  var dataRepresentation: Data { get }

  func sharedSecretFromKeyAgreement(with publicKeyShare: P256.KeyAgreement.PublicKey) throws
    -> SharedSecret
}

#if !os(Linux) && !os(Windows)
  class CryptoKitCrypto: Crypto {
    let context = LAContext()

    var isSecureEnclaveAvailable: Bool {
      return SecureEnclave.isAvailable
    }

    func newSecureEnclavePrivateKey(dataRepresentation: Data) throws -> SecureEnclavePrivateKey {
      return try SecureEnclave.P256.KeyAgreement.PrivateKey(
        dataRepresentation: dataRepresentation, authenticationContext: context)
    }

    func newSecureEnclavePrivateKey(accessControl: SecAccessControl) throws
      -> SecureEnclavePrivateKey
    {
      return try SecureEnclave.P256.KeyAgreement.PrivateKey(
        accessControl: accessControl, authenticationContext: context)
    }

    func newEphemeralPrivateKey() -> P256.KeyAgreement.PrivateKey {
      return P256.KeyAgreement.PrivateKey()
    }
  }

  extension SecureEnclave.P256.KeyAgreement.PrivateKey: SecureEnclavePrivateKey {
  }

#else

  class CryptoKitCrypto: Crypto {
    var isSecureEnclaveAvailable: Bool {
      return false
    }

    func newSecureEnclavePrivateKey(dataRepresentation: Data) throws -> SecureEnclavePrivateKey {
      throw Plugin.Error.seUnsupported
    }

    func newSecureEnclavePrivateKey(accessControl: SecAccessControl) throws
      -> SecureEnclavePrivateKey
    {
      throw Plugin.Error.seUnsupported
    }

    func newEphemeralPrivateKey() -> P256.KeyAgreement.PrivateKey {
      return P256.KeyAgreement.PrivateKey()
    }
  }

#endif

import Foundation

#if !os(Linux) && !os(Windows)
  import CryptoKit
#else
  import Crypto
#endif

class Plugin {
  var crypto: Crypto
  var stream: Stream

  init(crypto: Crypto, stream: Stream) {
    self.crypto = crypto
    self.stream = stream
  }

  func generateKey(accessControl: KeyAccessControl, recipientType: RecipientType, now: Date) throws
    -> (String, String)
  {
    if !crypto.isSecureEnclaveAvailable {
      throw Error.seUnsupported
    }
    #if !os(Linux) && !os(Windows)
      let createdAt = now.ISO8601Format()
      var accessControlFlags: SecAccessControlCreateFlags = [.privateKeyUsage]
      if accessControl == .anyBiometry || accessControl == .anyBiometryAndPasscode {
        accessControlFlags.insert(.biometryAny)
      }
      if accessControl == .currentBiometry || accessControl == .currentBiometryAndPasscode {
        accessControlFlags.insert(.biometryCurrentSet)
      }
      if accessControl == .passcode || accessControl == .anyBiometryAndPasscode
        || accessControl == .currentBiometryAndPasscode
      {
        accessControlFlags.insert(.devicePasscode)
      }
      if accessControl == .anyBiometryOrPasscode {
        accessControlFlags.insert(.userPresence)
      }
      var error: Unmanaged<CFError>?
      guard
        let secAccessControl = SecAccessControlCreateWithFlags(
          kCFAllocatorDefault, kSecAttrAccessibleWhenUnlockedThisDeviceOnly,
          accessControlFlags,
          &error)
      else {
        throw error!.takeRetainedValue() as Swift.Error
      }
    #else
      // FIXME: ISO8601Format currently not supported on Linux:
      //   https://github.com/apple/swift-corelibs-foundation/issues/4618
      // This code is only reached in unit tests on Linux anyway
      let createdAt = "1997-02-02T02:26:51Z"
      let secAccessControl = SecAccessControl()
    #endif

    let privateKey = try crypto.newSecureEnclavePrivateKey(accessControl: secAccessControl)
    let recipient = privateKey.publicKey.ageRecipient(type: recipientType)
    let identity = privateKey.ageIdentity
    let accessControlStr: String
    switch accessControl {
    case .none: accessControlStr = "none"
    case .passcode: accessControlStr = "passcode"
    case .anyBiometry: accessControlStr = "any biometry"
    case .anyBiometryOrPasscode: accessControlStr = "any biometry or passcode"
    case .anyBiometryAndPasscode: accessControlStr = "any biometry and passcode"
    case .currentBiometry: accessControlStr = "current biometry"
    case .currentBiometryAndPasscode: accessControlStr = "current biometry and passcode"
    }

    let contents = """
      # created: \(createdAt)
      # access control: \(accessControlStr)
      # public key: \(recipient)
      \(identity)
      """

    return (contents, recipient)
  }

  func generateRecipients(input: String, recipientType: RecipientType) throws -> String {
    var recipients: [String] = []
    for l in input.split(whereSeparator: \.isNewline) {
      if l.hasPrefix("#") {
        continue
      }
      let sl = String(l.trimmingCharacters(in: .whitespacesAndNewlines))
      let privateKey = try newSecureEnclavePrivateKey(ageIdentity: sl, crypto: self.crypto)
      recipients.append(privateKey.publicKey.ageRecipient(type: recipientType))
    }
    return recipients.joined(separator: "\n")
  }

  func runRecipientV1() {
    var recipients: [String] = []
    var identities: [String] = []
    var fileKeys: [Data] = []

    // Phase 1
    loop: while true {
      let stanza = try! Stanza.readFrom(stream: stream)
      switch stanza.type {
      case "add-recipient":
        recipients.append(stanza.args[0])
      case "add-identity":
        identities.append(stanza.args[0])
      case "wrap-file-key":
        fileKeys.append(stanza.body)
      case "done":
        break loop
      default:
        continue
      }
    }

    // Phase 2
    var stanzas: [Stanza] = []
    var errors: [Stanza] = []
    var recipientKeys: [(P256.KeyAgreement.PublicKey, RecipientStanzaType)] = []
    for (index, recipient) in recipients.enumerated() {
      do {
        recipientKeys.append(
          (
            (try P256.KeyAgreement.PublicKey(ageRecipient: recipient)),
            recipient.starts(with: "age1p256tag1") ? .p256tag : .pivp256
          ))
      } catch {
        errors.append(
          Stanza(error: "recipient", args: [String(index)], message: error.localizedDescription))
      }
    }
    for (index, identity) in identities.enumerated() {
      do {
        recipientKeys.append(
          (
            (try newSecureEnclavePrivateKey(ageIdentity: identity, crypto: crypto)).publicKey,
            .pivp256
          ))
      } catch {
        errors.append(
          Stanza(error: "identity", args: [String(index)], message: error.localizedDescription))
      }
    }
    for (index, fileKey) in fileKeys.enumerated() {
      for (recipientKey, recipientStanzaType) in recipientKeys {
        do {
          let ephemeralSecretKey = self.crypto.newEphemeralPrivateKey()
          let ephemeralPublicKeyBytes = ephemeralSecretKey.publicKey.compressedRepresentation
          // CryptoKit PublicKeys can be the identity point by construction (see CryptoTests), but
          // these keys can't be used in any operation. This is undocumented, but a documentation request
          // has been filed as FB11989432.
          // Swift Crypto PublicKeys cannot be the identity point by construction.
          // Compresed representation cannot be the identity point anyway (?)
          // Therefore, the shared secret cannot be all 0x00 bytes, so we don't need
          // to explicitly check this here.
          let sharedSecret = try ephemeralSecretKey.sharedSecretFromKeyAgreement(
            with: recipientKey)
          let sealedBox = try ChaChaPoly.seal(
            fileKey,
            using: recipientStanzaWrapKey(
              using: sharedSecret,
              salt: ephemeralPublicKeyBytes + recipientKey.compressedRepresentation,
              type: recipientStanzaType
            ), nonce: try! ChaChaPoly.Nonce(data: Data(count: 12)))
          stanzas.append(
            Stanza(
              type: "recipient-stanza",
              args: [
                String(index),
                recipientStanzaType.rawValue,
                recipientStanzaType == .p256tag
                  ? recipientKey.hmacTag(using: SymmetricKey(data: ephemeralPublicKeyBytes))
                    .base64RawEncodedString : recipientKey.sha256Tag.base64RawEncodedString,
                ephemeralPublicKeyBytes.base64RawEncodedString,
              ], body: sealedBox.ciphertext + sealedBox.tag
            )
          )
        } catch {
          errors.append(
            Stanza(error: "internal", args: [], message: error.localizedDescription))
        }
      }
    }
    for stanza in (errors.isEmpty ? stanzas : errors) {
      stanza.writeTo(stream: stream)
      let resp = try! Stanza.readFrom(stream: stream)
      assert(resp.type == "ok")
    }
    Stanza(type: "done").writeTo(stream: stream)
  }

  func runIdentityV1() {
    // Phase 1
    var identities: [String] = []
    var recipientStanzas: [Stanza] = []
    loop: while true {
      let stanza = try! Stanza.readFrom(stream: stream)
      switch stanza.type {
      case "add-identity":
        identities.append(stanza.args[0])
      case "recipient-stanza":
        recipientStanzas.append(stanza)
      case "done":
        break loop
      default:
        continue
      }
    }

    // Phase 2
    var identityKeys: [SecureEnclavePrivateKey] = []
    var errors: [Stanza] = []

    // Construct identities
    for (index, identity) in identities.enumerated() {
      do {
        identityKeys.append(
          (try newSecureEnclavePrivateKey(ageIdentity: identity, crypto: crypto)))
      } catch {
        errors.append(
          Stanza(error: "identity", args: [String(index)], message: error.localizedDescription))
      }
    }

    var fileResponses: [Int: Stanza] = [:]
    if errors.isEmpty {
      // Check structural validity
      for recipientStanza in recipientStanzas {
        let fileIndex = Int(recipientStanza.args[0])!
        switch recipientStanza.args[1] {
        case "piv-p256", "p256tag":
          if recipientStanza.args.count != 4 {
            fileResponses[fileIndex] = Stanza(
              error: "stanza", args: [String(fileIndex)], message: "incorrect argument count")
            continue
          }
          let tag = Data(base64RawEncoded: recipientStanza.args[2])
          if tag == nil || tag!.count != 4 {
            fileResponses[fileIndex] = Stanza(
              error: "stanza", args: [String(fileIndex)], message: "invalid tag")
            continue
          }
          let share = Data(base64RawEncoded: recipientStanza.args[3])
          if share == nil || share!.count != 33 {
            fileResponses[fileIndex] = Stanza(
              error: "stanza", args: [String(fileIndex)], message: "invalid share")
            continue
          }
          if recipientStanza.body.count != 32 {
            fileResponses[fileIndex] = Stanza(
              error: "stanza", args: [String(fileIndex)],
              message: "invalid body")
            continue
          }

        default:
          continue
        }
      }

      // Unwrap keys
      for recipientStanza in recipientStanzas {
        let fileIndex = Int(recipientStanza.args[0])!
        if fileResponses[fileIndex] != nil {
          continue
        }
        guard let type = RecipientStanzaType(rawValue: recipientStanza.args[1]) else {
          continue
        }
        let tag = recipientStanza.args[2]
        let share = recipientStanza.args[3]
        for identity in identityKeys {
          do {
            let shareKeyData = Data(base64RawEncoded: share)!
            let identityTag =
              type == .p256tag
              ? identity.publicKey.hmacTag(using: SymmetricKey(data: shareKeyData))
                .base64RawEncodedString : identity.publicKey.sha256Tag.base64RawEncodedString
            if identityTag != tag {
              continue
            }

            let shareKey: P256.KeyAgreement.PublicKey = try P256.KeyAgreement.PublicKey(
              compressedRepresentation: shareKeyData)
            // CryptoKit PublicKeys can be the identity point by construction (see CryptoTests), but
            // these keys can't be used in any operation. This is undocumented, but a documentation request
            // has been filed as FB11989432.
            // Swift Crypto PublicKeys cannot be the identity point by construction.
            // Compresed representation cannot be the identity point anyway (?)
            // Therefore, the shared secret cannot be all 0x00 bytes, so we don't need
            // to explicitly check this here.
            let sharedSecret: SharedSecret = try identity.sharedSecretFromKeyAgreement(
              with: shareKey)
            let unwrappedKey = try ChaChaPoly.open(
              ChaChaPoly.SealedBox(
                combined: try! ChaChaPoly.Nonce(data: Data(count: 12)) + recipientStanza.body),
              using: recipientStanzaWrapKey(
                using: sharedSecret,
                salt: shareKey.compressedRepresentation
                  + identity.publicKey.compressedRepresentation,
                type: type
              ))
            fileResponses[fileIndex] = Stanza(
              type: "file-key",
              args: [String(fileIndex)],
              body: unwrappedKey
            )
          } catch {
            Stanza(type: "msg", body: Data(error.localizedDescription.utf8)).writeTo(
              stream: stream)
            let resp = try! Stanza.readFrom(stream: self.stream)
            assert(resp.type == "ok")
            // continue
          }
        }
      }
    }

    let responses = fileResponses.keys.sorted().map({ k in fileResponses[k]! })
    for stanza in (errors.isEmpty ? responses : errors) {
      stanza.writeTo(stream: stream)
      let resp = try! Stanza.readFrom(stream: stream)
      assert(resp.type == "ok")
    }
    Stanza(type: "done").writeTo(stream: stream)
  }

  enum Error: LocalizedError, Equatable {
    case seUnsupported
    case incompleteStanza
    case invalidStanza
    case unknownHRP(String)

    public var errorDescription: String? {
      switch self {
      case .seUnsupported: return "Secure Enclave not supported on this device"
      case .incompleteStanza: return "incomplete stanza"
      case .invalidStanza: return "invalid stanza"
      case .unknownHRP(let hrp): return "unknown HRP: \(hrp)"
      }
    }
  }
}

//////////////////////////////////////////////////////////////////////////////////////////

struct Stanza: Equatable {
  var type: String
  var args: [String] = []
  var body = Data()

  static func readFrom(stream: Stream) throws -> Stanza {
    guard let header = stream.readLine() else {
      throw Plugin.Error.incompleteStanza
    }
    let headerParts = header.components(separatedBy: " ")
    if headerParts.count < 2 {
      throw Plugin.Error.invalidStanza
    }
    if headerParts[0] != "->" {
      throw Plugin.Error.invalidStanza
    }
    var body = Data()
    while true {
      guard let line = stream.readLine() else {
        throw Plugin.Error.incompleteStanza
      }
      guard let lineData = Data(base64RawEncoded: line) else {
        throw Plugin.Error.invalidStanza
      }
      if lineData.count > 48 {
        throw Plugin.Error.invalidStanza
      }
      body.append(lineData)
      if lineData.count < 48 {
        break
      }
    }
    return Stanza(type: headerParts[1], args: Array(headerParts[2...]), body: body)
  }

  func writeTo(stream: Stream) {
    let parts = ([type] + args).joined(separator: " ")
    stream.writeLine("-> \(parts)\n\(body.base64RawEncodedString)")
  }
}

extension Stanza {
  init(error type: String, args: [String] = [], message: String) {
    self.type = "error"
    self.args = [type] + args
    self.body = Data(message.utf8)
  }
}

enum KeyAccessControl {
  case none
  case passcode
  case anyBiometry
  case anyBiometryOrPasscode
  case anyBiometryAndPasscode
  case currentBiometry
  case currentBiometryAndPasscode
}

enum RecipientType: String {
  case se = "se"
  case p256tag = "p256tag"
}

enum RecipientStanzaType: String {
  case p256tag = "p256tag"
  case pivp256 = "piv-p256"
}

func recipientStanzaWrapKey(
  using sharedSecret: SharedSecret, salt: Data, type: RecipientStanzaType
) -> SymmetricKey {
  switch type {
  case .p256tag:
    return sharedSecret.hkdfDerivedSymmetricKey(
      using: SHA256.self,
      salt: Data("age-encryption.org/v1/p256tag".utf8),
      sharedInfo: salt,
      outputByteCount: 32
    )
  case .pivp256:
    return sharedSecret.hkdfDerivedSymmetricKey(
      using: SHA256.self, salt: salt,
      sharedInfo: Data("piv-p256".utf8),
      outputByteCount: 32
    )
  }
}

extension P256.KeyAgreement.PublicKey {
  init(ageRecipient: String) throws {
    let id = try Bech32().decode(ageRecipient)
    if id.hrp != "age1se" && id.hrp != "age1p256tag" {
      throw Plugin.Error.unknownHRP(id.hrp)
    }
    self = try P256.KeyAgreement.PublicKey(compressedRepresentation: id.data)
  }

  var sha256Tag: Data {
    return Data(SHA256.hash(data: compressedRepresentation).prefix(4))
  }

  func hmacTag(using: SymmetricKey) -> Data {
    return Data(
      HMAC<SHA256>.authenticationCode(for: compressedRepresentation, using: using).prefix(4))
  }

  func ageRecipient(type: RecipientType) -> String {
    return Bech32().encode(hrp: "age1\(type.rawValue)", data: self.compressedRepresentation)
  }
}

extension SecureEnclavePrivateKey {
  var ageIdentity: String {
    return Bech32().encode(
      hrp: "AGE-PLUGIN-SE-",
      data: self.dataRepresentation)
  }
}

func newSecureEnclavePrivateKey(ageIdentity: String, crypto: Crypto) throws
  -> SecureEnclavePrivateKey
{
  let id = try Bech32().decode(ageIdentity)
  if id.hrp != "AGE-PLUGIN-SE-" {
    throw Plugin.Error.unknownHRP(id.hrp)
  }
  return try crypto.newSecureEnclavePrivateKey(dataRepresentation: id.data)
}

import Foundation

/// Abstraction of a line-based communication stream
protocol Stream {
  func readLine() -> String?
  func writeLine(_: String)
}

class StandardIOStream: Stream {
  func readLine() -> String? {
    return Swift.readLine(strippingNewline: true)
  }

  func writeLine(_ line: String) {
    FileHandle.standardOutput.write(Data(line.utf8))
    FileHandle.standardOutput.write(Data([0xa]))
    fflush(stdout)
  }
}