usbtool.command

#!/usr/bin/swift

/*
 *  File: usbtool.command
 *
 *  Copyright © 2021-2022 by vulgo 74738713+vulgo@users.noreply.github.com
 *
 *  Permission to use, copy, modify, and/or distribute this software for any purpose with or without fee is hereby granted.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE
 *  INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 *  SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE,
 *  DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 *  CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

import Cocoa

public enum USBConnector: UInt8 {
	case a2 = 0x0 // Type A (USB 2)
	case a3 = 0x3 // Type A (USB 3)
	case c3s = 0x9 // Type C (USB 3, switched)
	case c3 = 0xA // Type C (USB 3)
	case proprietary = 0xFF // e.g. internal bluetooth
}

public protocol USBPort: AnyObject {
	var name: String { get }
	var address: UInt32 { get }
	var connector: USBConnector { get }
	var info: String { get }
	var isEnabled: Bool { get }
}

public protocol USBPortMap: AnyObject {
	var data: [USBPort] { get }
	var lastAddress: UInt32? { get }
	var enabledCount: Int { get }
}

public protocol USBInfoPlist: AnyObject {
	var dictionary: [String: Any] { get }
	var data: Data { get }
}

public protocol USBKextBundle: AnyObject {
	var bundleURL: URL { get }
	var contentsURL: URL { get }
	var infoPlistURL: URL  { get }
	var infoPlist: USBInfoPlist { get }
	var exists: Bool { get }
	func remove() throws
	func write() throws
}

let kTargetBoard = "B460i-PLUS"
let kAppTitle = "\(kTargetBoard) USB Tool"
let kBundleName = "USBPorts"
let kBundleIdentifier = "org.usbtool.kext"
let kDriverPersonalityKey = "Generated by \(kAppTitle)"
let kDriverIONameMatch = "XHC"
let kDriverIOProviderClass = "AppleUSBXHCISPT"
let kWriteButtonTitle = "Write Bundle to Desktop"
let kWriteSuccessFormat = "Wrote bundle to %@"

let kDefaultMapData: [Port] = [
	.init(
		name: "HS01",
		address: 0x1,
		connector: .a3,
		info: "USB 3 header"
	), .init(
		name: "HS02",
		address: 0x2,
		connector: .a3,
		info: "USB 3 header"
	), .init(
		name: "HS03",
		address: 0x3,
		connector: .a3,
		info: "Rear panel USB 3 connector"
	), .init(
		name: "HS04",
		address: 0x4,
		connector: .a3,
		info: "Rear panel USB 3 connector"
	), .init(
		name: "HS05",
		address: 0x5,
		connector: .a3,
		info: "Rear panel USB 3 connector"
	), .init(
		name: "HS06",
		address: 0x6,
		connector: .a3,
		info: "Rear panel USB 3 connector"
	), .init(
		name: "HS07",
		address: 0x7,
		connector: .a2,
		info: "Rear panel USB 2 connector"
	), .init(
		name: "HS08",
		address: 0x8,
		connector: .a2,
		info: "Rear panel USB 2 connector"
	), .init(
		name: "HS09",
		address: 0x9,
		connector: .a2,
		info: "USB 2 header"
	), .init(
		name: "HS10",
		address: 0xA,
		connector: .a2,
		info: "USB 2 header"
	), .init(
		name: "HS12",
		address: 0xC,
		connector: .proprietary,
		info: "m.2 E key"
	), .init(
		name: "SS01",
		address: 0x11,
		connector: .a3,
		info: "USB 3 header"
	), .init(
		name: "SS02",
		address: 0x12,
		connector: .a3,
		info: "USB 3 header"
	), .init(
		name: "SS03",
		address: 0x13,
		connector: .a3,
		info: "Rear panel USB 3 connector"
	), .init(
		name: "SS04",
		address: 0x14,
		connector: .a3,
		info: "Rear panel USB 3 connector"
	), .init(
		name: "SS05",
		address: 0x15,
		connector: .a3,
		info: "Rear panel USB 3 connector"
	), .init(
		name: "SS06",
		address: 0x16,
		connector: .a3,
		info: "Rear panel USB 3 connector"
	)
]

let kWindowWidth: CGFloat = 360.0
let kContentSpacing: CGFloat = 20.0
let kPortListEdgeInsets = NSEdgeInsets(top: 5.0, left: 10.0, bottom: 5.0, right: 10.0)
let kPortListColumnSpacing: CGFloat = 3.0
let kPortListRowSpacing: CGFloat = 9.0
let kPortListViewLeading = kContentSpacing + kPortListEdgeInsets.left
let kPortListViewTrailing = -(kContentSpacing + kPortListEdgeInsets.right)
let kPortListViewTop = kContentSpacing + kPortListEdgeInsets.top
let kPortListViewBottom = -(kContentSpacing + kPortListEdgeInsets.bottom)
let kPortLimit = 15
let kFileManager = FileManager.default

let kInfoPlistSourceDictionary: [String: Any] = {
	let sourceData = """
	PD94bWwgdmVyc2lvbj0iMS4wIiBlbmNvZGluZz0iVVRGLTgiPz48IURPQ1RZUEUgcGxpc3QgUFVCTElDICItLy
	9BcHBsZS8vRFREIFBMSVNUIDEuMC8vRU4iICJodHRwOi8vd3d3LmFwcGxlLmNvbS9EVERzL1Byb3BlcnR5TGlz
	dC0xLjAuZHRkIj48cGxpc3QgdmVyc2lvbj0iMS4wIj48ZGljdD48a2V5PkNGQnVuZGxlSW5mb0RpY3Rpb25hcn
	lWZXJzaW9uPC9rZXk+PHN0cmluZz42LjA8L3N0cmluZz48a2V5PkNGQnVuZGxlUGFja2FnZVR5cGU8L2tleT48
	c3RyaW5nPktFWFQ8L3N0cmluZz48a2V5PkNGQnVuZGxlVmVyc2lvbjwva2V5PjxzdHJpbmc+MS4wPC9zdHJpbm
	c+PGtleT5PU0J1bmRsZVJlcXVpcmVkPC9rZXk+PHN0cmluZz5Sb290PC9zdHJpbmc+PC9kaWN0PjwvcGxpc3Q+
	"""
	
	do {
		guard let data = Data(base64Encoded: sourceData, options: [.ignoreUnknownCharacters]) else {
			throw USBTool.RuntimeError("failed to decode property list data from base64")
		}
		
		let plist: Any = try PropertyListSerialization.propertyList(from: data, options: .mutableContainersAndLeaves, format: nil)
		
		guard let result = plist as? [String: Any] else {
			throw USBTool.RuntimeError("failed to cast serialized property list to dictionary")
		}
		
		return result
	}
	
	catch {
		fatalError(error.localizedDescription)
	}
}()

let kModelIdentifier: String = {
	var identifier: String?
	let entry = IORegistryEntryFromPath(kIOMainPortDefault, "IOService:/")
	
	if let data = IORegistryEntryCreateCFProperty(entry, "model" as CFString, kCFAllocatorDefault, 0)?.takeRetainedValue() as? Data {
		identifier = String(data: data, encoding: .macOSRoman)?.filter { $0 != "\0"}
	}
	
	IOObjectRelease(entry)
	
	guard identifier != nil else {
		fatalError("failed to obtain a 'model' property from the IORegistry")
	}
	
	return identifier!
}()

#if compiler(<5.5)
let kIOMainPortDefault: mach_port_t = kIOMasterPortDefault
#endif

public extension DispatchQueue {
	static let work = DispatchQueue(label: "work")
}

extension FileHandle: TextOutputStream {
	public func write(_ string: String) {
		write(string.data(using: .utf8)!)
	}
}

public extension FileManager {
	func directoryExists(atPath path: String) -> Bool {
		var dir = ObjCBool(false)
		return fileExists(atPath: path, isDirectory: &dir) && dir.boolValue
	}
	
	func itemExists(atPath path: String) -> Bool {
		var dir = ObjCBool(false)
		return fileExists(atPath: path, isDirectory: &dir)
	}
	
	func desktopDirectoryURL() throws -> URL {
		guard let url = urls(for: .desktopDirectory, in: .userDomainMask).first else {
			throw USBTool.RuntimeError("failed to obtain desktop directory URL")
		}
		
		return url
	}
}

public extension FixedWidthInteger {
	var data: Data {
		var this = self
		return .init(bytes: &this, count: MemoryLayout<Self>.size)
	}
}

public extension NSMenuItem {
	convenience init(title: String, action: Selector? = nil) {
		self.init(title: title, action: action, keyEquivalent: "")
	}
}

public extension NSMenu {
	convenience init(title: String, items: [NSMenuItem]) {
		self.init(title: title)
		self.items = items
	}
}

public extension NSWindow {
	var closeButton: NSButton? {
		standardWindowButton(.closeButton)
	}
	
	var miniaturizeButton: NSButton? {
		standardWindowButton(.miniaturizeButton)
	}
	
	var zoomButton: NSButton? {
		standardWindowButton(.zoomButton)
	}
}

public extension PropertyListSerialization {
	class func xmlData(from dictionary: [String: Any]) throws -> Data {
		return try Self.data(fromPropertyList: dictionary, format: .xml, options: 0)
	}
}

public extension NSTextField {
	static func makePortLabel(_ stringValue: String) -> NSTextField {
		let textField = NSTextField()
		textField.drawsBackground = false
		textField.isBezeled = false
		textField.isEditable = false
		textField.isSelectable = true
		textField.textColor = NSColor.secondaryLabelColor
		textField.stringValue = stringValue
		textField.invalidateIntrinsicContentSize()
		return textField
	}
}

public extension NSButton {
	static func makePortSwitchButton(title: String, enabled: Bool) -> NSButton {
		let button = NSButton()
		button.title = title
		button.setButtonType(.switch)
		button.state = enabled ? .on : .off
		button.invalidateIntrinsicContentSize()
		return button
	}
}

public extension NSAlert {
	convenience init(_ messageText: String, buttonTitles: String ...) {
		self.init()
		self.messageText = messageText
		
		guard !buttonTitles.isEmpty else {
			addButton(withTitle: "Close")
			return
		}
		
		for title in buttonTitles {
			addButton(withTitle: title)
		}
	}
}

public extension USBPortMap {
	var lastAddress: UInt32? {
		return data.sorted { $0.address < $1.address }.last?.address
	}
	
	var enabledCount: Int {
		return data.filter { $0.isEnabled }.count
	}
}

public extension USBKextBundle {
	var exists: Bool {
		kFileManager.itemExists(atPath: bundleURL.path)
	}
	
	func remove() throws {
		try kFileManager.removeItem(atPath: bundleURL.path)
	}
	
	func write() throws {
		if exists {
			try DispatchQueue.main.sync {
				if NSAlert("\(bundleURL.path) exists", buttonTitles: "Overwrite", "Cancel").runModal() != .alertFirstButtonReturn {
					throw USBTool.RuntimeError("cancelled by user")
				}
				
				try remove()
			}
		}
		
		try kFileManager.createDirectory(atPath: contentsURL.path, withIntermediateDirectories: true, attributes: nil)
		try infoPlist.data.write(to: infoPlistURL, options: .atomic)
	}
}

public func /(url: URL, pathComponent: String) -> URL {
	return url.appendingPathComponent(pathComponent)
}

internal final class Port: NSObject, USBPort {
	public let name: String
	public let address: UInt32
	public let connector: USBConnector
	public let info: String
	@objc private(set) var isEnabled: Bool
	
	public init(name: String, address: UInt32, connector: USBConnector, info: String, isEnabled: Bool = true) {
		self.name = name
		self.address = address
		self.connector = connector
		self.info = info
		self.isEnabled = isEnabled
	}
}

internal final class PortMap: USBPortMap {
	static let `default` = PortMap(data: kDefaultMapData)
	public let data: [USBPort]
	
	public init(data: [USBPort]) {
		self.data = data
	}
}

internal final class InfoPlist: USBInfoPlist {
	private static let bundleIdentifierKey = kCFBundleIdentifierKey as String
	private static let bundleNameKey = kCFBundleNameKey as String
	private static let driverBundleIdentifier = "com.apple.driver.AppleUSBHostMergeProperties"
	private static let driverClass = "AppleUSBHostMergeProperties"
	private static let personalitiesKey = "IOKitPersonalities"
	private static let providerMergePropertiesKey = "IOProviderMergeProperties"
	private static let modelIdentifierKey = "model"
	private static let portCountKey = "port-count"
	private static let portsKey = "ports"
	private static let portAddressKey = "port"
	private static let portConnectorKey = "UsbConnector"
	private(set) var dictionary: [String: Any]
	private(set) var data: Data
	
	public init(portMap: USBPortMap) throws {
		dictionary = kInfoPlistSourceDictionary
		data = Data()
		
		guard let portCount = portMap.lastAddress else {
			throw USBTool.RuntimeError("the port map is empty")
		}
		
		var ports = [String: Any]()
		
		for port in portMap.data {
			ports[port.name] = [
				Self.portAddressKey: port.address.data,
				Self.portConnectorKey: port.connector.rawValue
			] as [String: Any]
		}
		
		let driverPersonality: [String: Any] = [
			Self.bundleIdentifierKey: Self.driverBundleIdentifier,
			kIOClassKey: Self.driverClass,
			kIONameMatchKey: kDriverIONameMatch,
			kIOProviderClassKey: kDriverIOProviderClass,
			Self.modelIdentifierKey: kModelIdentifier,
			Self.providerMergePropertiesKey: [
				Self.portCountKey: portCount.data,
				Self.portsKey: ports
			] as [String: Any]
		]
		
		dictionary[Self.bundleNameKey] = kBundleName
		dictionary[Self.bundleIdentifierKey] = kBundleIdentifier
		dictionary[Self.personalitiesKey] = [
			kDriverPersonalityKey: driverPersonality
		]
		
		data = try PropertyListSerialization.xmlData(from: dictionary)
	}
}

internal final class KextBundle: USBKextBundle {
	public let bundleURL: URL
	public let contentsURL: URL
	public let infoPlistURL: URL
	public let infoPlist: USBInfoPlist
	
	public init(location: URL, infoPlist: USBInfoPlist) throws {
		guard kFileManager.directoryExists(atPath: location.path) else {
			throw USBTool.RuntimeError("directory not found at \(location.path)")
		}
		
		bundleURL = location / "\(kBundleName).kext"
		contentsURL = bundleURL / "Contents"
		infoPlistURL = contentsURL / "Info.plist"
		self.infoPlist = infoPlist
	}
}

internal final class ViewController: NSViewController {
	public static let shared = ViewController()
	
	private init() {
		super.init(nibName: nil, bundle: nil)
	}
	
	required init?(coder: NSCoder) {
		return nil
	}
	
	private lazy var portListView: NSGridView = {
		let view = NSGridView(numberOfColumns: 2, rows: 0)
		
		for port in PortMap.default.data {
			let button = NSButton.makePortSwitchButton(title: port.name, enabled: port.isEnabled)
			let label = NSTextField.makePortLabel(port.info)
			button.target = self
			button.action = #selector(Self.switchButtonPressed(_:))
			button.bind(.value, to: port,
				    withKeyPath: #keyPath(Port.isEnabled),
				    options: [.validatesImmediately: true])
			view.addRow(with: [button, label])
		}
		
		view.rowAlignment = .firstBaseline
		view.columnSpacing = kPortListColumnSpacing
		view.rowSpacing = kPortListRowSpacing
		view.column(at: 0).xPlacement = .fill
		view.column(at: 1).xPlacement = .leading
		view.translatesAutoresizingMaskIntoConstraints = false
		return view
	}()
	
	private lazy var writeButton: NSButton = {
		let button = NSButton()
		button.bezelStyle = .rounded
		button.title = kWriteButtonTitle
		button.isEnabled = false
		button.target = self
		button.action = #selector(Self.writeButtonPressed(_:))
		button.translatesAutoresizingMaskIntoConstraints = false
		return button
	}()
	
	override func loadView() {
		view = USBTool.mainWindow.contentView!
	}
	
	private func activateLayoutConstraints() {
		NSLayoutConstraint.activate([
			view.widthAnchor.constraint(greaterThanOrEqualToConstant: kWindowWidth),
			portListView.leadingAnchor.constraint(equalTo: view.leadingAnchor, constant: kPortListViewLeading),
			portListView.trailingAnchor.constraint(equalTo: view.trailingAnchor, constant: kPortListViewTrailing),
			portListView.topAnchor.constraint(equalTo: view.topAnchor, constant: kPortListViewTop),
			portListView.bottomAnchor.constraint(equalTo: writeButton.topAnchor, constant: kPortListViewBottom),
			writeButton.bottomAnchor.constraint(equalTo: view.bottomAnchor, constant: -kContentSpacing),
			writeButton.trailingAnchor.constraint(equalTo: view.trailingAnchor, constant: -kContentSpacing)
		])
	}
	
	override func viewDidLoad() {
		view.subviews = [portListView, writeButton]
		activateLayoutConstraints()
		enableWriteButton()
	}
	
	private func enableWriteButton() {
		writeButton.isEnabled = PortMap.default.enabledCount <= kPortLimit
	}
	
	@objc func switchButtonPressed(_ sender: NSButton) {
		enableWriteButton()
	}
	
	@objc func writeButtonPressed(_ sender: NSButton) {
		let userMap = PortMap(data: PortMap.default.data.filter { $0.isEnabled })
		
		DispatchQueue.work.async { [userMap] in
			do {
				let destinationURL = try kFileManager.desktopDirectoryURL()
				var path: String = .init()
				try USBTool.writeBundle(destination: destinationURL, portMap: userMap, bundlePath: &path)
				let message = String(format: kWriteSuccessFormat, path)
				
				DispatchQueue.main.async {
					if NSAlert(message, buttonTitles: "Close", "Reveal in Finder").runModal() == .alertSecondButtonReturn {
						NSWorkspace.shared.selectFile(path, inFileViewerRootedAtPath: destinationURL.path)
					}
				}
			}
			
			catch {
				var errorStream = FileHandle.standardError
				print(error.localizedDescription, to: &errorStream)
			}
		}
	}
}

internal final class AppDelegate: NSObject, NSApplicationDelegate {
	public static let shared = AppDelegate()
	
	private override init() {
		super.init()
	}
	
	@objc func applicationShouldTerminateAfterLastWindowClosed(_ sender: NSApplication) -> Bool {
		return true
	}
	
	@objc func applicationDidFinishLaunching(_ notification: Notification) {
		NSApp.activate(ignoringOtherApps: true)
	}
}

internal struct USBTool {
	struct RuntimeError: LocalizedError {
		let description: String
		let location: String
		
		init(_ description: String, location: String = #function) {
			self.description = description
			self.location = location
		}
		
		var errorDescription: String? {
			return "The operation couldn't be completed: \(String(describing: self))"
		}
	}
	
	static let mainMenu: NSMenu = {
		let processName = ProcessInfo.processInfo.processName
		let appMenu = NSMenuItem(title: processName)
		let editMenu = NSMenuItem(title: "Edit")
		appMenu.submenu = NSMenu(title: processName, items: [
			NSMenuItem(title: "About \(processName)",
				   action: #selector(NSApplication.orderFrontStandardAboutPanel(_:))),
			NSMenuItem.separator(),
			NSMenuItem(title: "Quit \(processName)",
				   action: #selector(NSApplication.terminate(_:)),
				   keyEquivalent: "q")
		])
		editMenu.submenu = NSMenu(title: "Edit", items: [
			NSMenuItem(title: "Copy",
				   action: #selector(NSText.copy(_:)),
				   keyEquivalent: "c")
		])
		return NSMenu(title: "Main Menu", items: [appMenu, editMenu])
	}()
	
	static let mainWindow: NSWindow = {
		let window = NSWindow()
		window.title = kAppTitle
		window.styleMask = [.titled, .closable]
		window.miniaturizeButton?.isHidden = true
		window.zoomButton?.isHidden = true
		return window
	}()
	
	static func writeBundle(destination url: URL, portMap: USBPortMap, bundlePath: inout String) throws {
		let kext = try KextBundle(location: url, infoPlist: try InfoPlist(portMap: portMap))
		try kext.write()
		bundlePath = kext.bundleURL.path
	}
	
	static func main() {
		let window = Self.mainWindow
		NSApp = NSApplication.shared
		NSApp.delegate = AppDelegate.shared
		NSApp.mainMenu = Self.mainMenu
		NSApp.setActivationPolicy(.regular)
		window.contentViewController = ViewController.shared
		window.makeKeyAndOrderFront(self)
		window.center()
		NSApp.run()
	}
}

USBTool.main()