1/*
2 * Copyright (c) 2002 Apple Inc. All rights reserved.
3 *
4 * @APPLE_LICENSE_HEADER_START@
5 *
6 * This file contains Original Code and/or Modifications of Original Code
7 * as defined in and that are subject to the Apple Public Source License
8 * Version 2.0 (the 'License'). You may not use this file except in
9 * compliance with the License. Please obtain a copy of the License at
10 * http://www.opensource.apple.com/apsl/ and read it before using this
11 * file.
12 *
13 * The Original Code and all software distributed under the License are
14 * distributed on an 'AS IS' basis, WITHOUT WARRANTY OF ANY KIND, EITHER
15 * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES,
16 * INCLUDING WITHOUT LIMITATION, ANY WARRANTIES OF MERCHANTABILITY,
17 * FITNESS FOR A PARTICULAR PURPOSE, QUIET ENJOYMENT OR NON-INFRINGEMENT.
18 * Please see the License for the specific language governing rights and
19 * limitations under the License.
20 *
21 * @APPLE_LICENSE_HEADER_END@
22 */
23
24/*
25 * bsdpc.c
26 * - command line utility for selecting a netboot image
27 */
28
29/*
30 * Modification History
31 *
32 * February 25, 2002	Dieter Siegmund (dieter@apple.com)
33 * - initial revision
34 */
35
36#include <stdlib.h>
37#include <stdio.h>
38#include <sys/types.h>
39#include <sys/errno.h>
40#include <sys/socket.h>
41#include <sys/ioctl.h>
42#include <sys/fcntl.h>
43#include <mach/boolean.h>
44#include <limits.h>
45
46#include "util.h"
47#include "cfutil.h"
48#include "BSDPClient.h"
49#include "BSDPClientPrivate.h"
50#include <CoreFoundation/CoreFoundation.h>
51
52static __inline__ const char *
53image_kind_string(int kind)
54{
55    const char * kind_str[] = {
56	"Mac OS 9",
57	"Mac OS X",
58	"Mac OS X Server",
59    };
60    if (kind >= kBSDPImageKindMacOS9
61	&& kind <= kBSDPImageKindMacOSXServer) {
62	return (kind_str[kind]);
63    }
64    return (NULL);
65}
66
67
68#define kServerAddress		CFSTR("_bsdpc_ServerAddress")
69#define kServerPriority		CFSTR("_bsdpc_ServerPriority")
70#define kImageList		CFSTR("_bsdpc_ImageList")
71#define kServedBy		CFSTR("_bsdpc_ServedBy")
72
73typedef enum {
74    bsdpc_state_init = 0,
75    bsdpc_state_list,
76    bsdpc_state_user_input,
77    bsdpc_state_select,
78} bsdpc_state_t;
79
80#define INVALID_REDISPLAY	5
81
82struct bsdpc;
83typedef struct bsdpc bsdpc_t;
84
85typedef void (*timerCallBack)(bsdpc_t * bsdpc);
86
87struct bsdpc {
88    BSDPClientRef		client;
89    bsdpc_state_t		state;
90    boolean_t			gathering;
91    int				invalid_count;
92    CFMutableArrayRef		images;
93    CFMutableArrayRef		menu;
94    CFRunLoopTimerRef		timer;
95    timerCallBack		timer_callback;
96};
97
98static bsdpc_t	bsdpc;
99
100static void
101processTimer(CFRunLoopTimerRef timer, void * info)
102{
103    bsdpc_t *	bsdpc;
104
105    bsdpc = (bsdpc_t *)info;
106    (*bsdpc->timer_callback)(bsdpc);
107    return;
108}
109
110static void
111cancelTimer(bsdpc_t * bsdpc)
112{
113    if (bsdpc->timer) {
114	CFRunLoopTimerInvalidate(bsdpc->timer);
115	my_CFRelease(&bsdpc->timer);
116    }
117    bsdpc->timer_callback = NULL;
118    return;
119}
120
121static void
122setTimer(bsdpc_t * bsdpc, struct timeval rel_time,
123	 timerCallBack callback)
124{
125    CFRunLoopTimerContext 	context =  { 0, NULL, NULL, NULL, NULL };
126    CFAbsoluteTime 		wakeup_time;
127
128    cancelTimer(bsdpc);
129    bsdpc->timer_callback = callback;
130    wakeup_time = CFAbsoluteTimeGetCurrent() + rel_time.tv_sec
131	  + ((double)rel_time.tv_usec / USECS_PER_SEC);
132    context.info = bsdpc;
133    bsdpc->timer
134	= CFRunLoopTimerCreate(NULL, wakeup_time,
135			       0.0, 0, 0,
136			       processTimer,
137			       &context);
138    CFRunLoopAddTimer(CFRunLoopGetCurrent(), bsdpc->timer,
139		      kCFRunLoopDefaultMode);
140    return;
141}
142
143static void
144add_menu_item(CFMutableArrayRef menu, CFDictionaryRef server_dict,
145	      CFDictionaryRef image_dict)
146{
147    int				i;
148    CFNumberRef			identifier;
149    CFMutableArrayRef		served_by;
150    CFMutableDictionaryRef	menu_item;
151
152    identifier = CFDictionaryGetValue(image_dict,
153				      kBSDPImageDescriptionIdentifier);
154    if (BSDPImageDescriptionIdentifierIsServerLocal(identifier) == FALSE) {
155	int count = CFArrayGetCount(menu);
156	for (i = 0; i < count; i++) {
157	    menu_item = (CFMutableDictionaryRef)CFArrayGetValueAtIndex(menu, i);
158	    if (CFEqual(identifier,
159			CFDictionaryGetValue(menu_item,
160					     kBSDPImageDescriptionIdentifier))
161		== TRUE) {
162		served_by = (CFMutableArrayRef)
163		    CFDictionaryGetValue(menu_item, kServedBy);
164		CFArrayAppendValue(served_by, server_dict);
165		return;
166	    }
167	}
168    }
169    menu_item = CFDictionaryCreateMutableCopy(NULL, 0, image_dict);
170    served_by = CFArrayCreateMutable(NULL, 0, &kCFTypeArrayCallBacks);
171    CFArrayAppendValue(served_by, server_dict);
172    CFDictionarySetValue(menu_item, kServedBy, served_by);
173    CFArrayAppendValue(menu, menu_item);
174    my_CFRelease(&served_by);
175    my_CFRelease(&menu_item);
176    return;
177}
178
179static void
180create_menu(bsdpc_t * bsdpc)
181{
182    int count;
183    int i;
184    int	s;
185
186    my_CFRelease(&bsdpc->menu);
187    bsdpc->menu = CFArrayCreateMutable(NULL, 0, &kCFTypeArrayCallBacks);
188    count = CFArrayGetCount(bsdpc->images);
189    for (s = 0; s < count; s++) {
190	CFArrayRef		images;
191	int			images_count;
192	CFMutableDictionaryRef	server_copy;
193	CFDictionaryRef		server_dict;
194
195	server_dict = CFArrayGetValueAtIndex(bsdpc->images, s);
196	server_copy = CFDictionaryCreateMutableCopy(NULL, 0, server_dict);
197	CFDictionaryRemoveValue(server_copy, kImageList);
198	images = CFDictionaryGetValue(server_dict, kImageList);
199	images_count = CFArrayGetCount(images);
200	for (i = 0; i < images_count; i++) {
201	    CFDictionaryRef		image_dict;
202
203	    image_dict = CFArrayGetValueAtIndex(images, i);
204	    add_menu_item(bsdpc->menu, server_copy, image_dict);
205	}
206	my_CFRelease(&server_copy);
207    }
208    return;
209}
210
211static int
212cfstring_to_cstring(CFStringRef cfstr, char * str, int len)
213{
214    CFIndex		l;
215    CFRange		range;
216
217    range = CFRangeMake(0, CFStringGetLength(cfstr));
218    (void)CFStringGetBytes(cfstr, range, kCFStringEncodingUTF8,
219			   '?', TRUE, (UInt8 *)str, len, &l);
220    str[l] = '\0';
221    return (l);
222}
223
224static void
225display_prompt(bsdpc_t * bsdpc)
226{
227    printf("Enter the image number (1..%ld) or q to quit) ",
228	   CFArrayGetCount(bsdpc->menu));
229    fflush(stdout);
230    return;
231}
232
233static void
234display_menu(bsdpc_t * bsdpc)
235{
236    int		count;
237    int		i;
238
239    count = CFArrayGetCount(bsdpc->menu);
240    printf("\nNetBoot Image List:\n");
241    for (i = 0; i < count; i++) {
242	BSDPImageKind		kind;
243	const char *		kind_str;
244	CFDictionaryRef		menu_item;
245	char			name[256];
246	CFBooleanRef		is_default;
247	CFNumberRef		identifier;
248	CFBooleanRef		is_selected;
249
250	menu_item = CFArrayGetValueAtIndex(bsdpc->menu, i);
251	is_selected = CFDictionaryGetValue(menu_item,
252					   kBSDPImageDescriptionIsSelected);
253	is_default = CFDictionaryGetValue(menu_item,
254					   kBSDPImageDescriptionIsDefault);
255	identifier = CFDictionaryGetValue(menu_item,
256					  kBSDPImageDescriptionIdentifier);
257	cfstring_to_cstring(CFDictionaryGetValue(menu_item,
258						 kBSDPImageDescriptionName),
259			    name, sizeof(name));
260	printf("%4d. %s", i + 1, name);
261
262	kind = BSDPImageDescriptionIdentifierImageKind(identifier);
263	kind_str = image_kind_string(kind);
264	if (kind_str != NULL) {
265	    printf(" [%s]", kind_str);
266	}
267	else {
268	    printf(" [Kind=%d]", kind);
269	}
270	if (BSDPImageDescriptionIdentifierIsInstall(identifier)) {
271	    printf(" [Install]");
272	}
273	if (is_default != NULL
274	    && CFEqual(is_default, kCFBooleanTrue)) {
275	    printf(" [Default]");
276	}
277	if (is_selected != NULL
278	    && CFEqual(is_selected, kCFBooleanTrue)) {
279	    printf(" [Selected]");
280	}
281	printf("\n");
282    }
283    display_prompt(bsdpc);
284    return;
285}
286
287static void
288redisplay(bsdpc_t * bsdpc)
289{
290    bsdpc->invalid_count++;
291    if (bsdpc->invalid_count == INVALID_REDISPLAY) {
292	bsdpc->invalid_count = 0;
293	display_menu(bsdpc);
294    }
295    else {
296	display_prompt(bsdpc);
297    }
298}
299
300static void
301gatherDone(bsdpc_t * bsdpc)
302{
303    bsdpc->state = bsdpc_state_user_input;
304    create_menu(bsdpc);
305    display_menu(bsdpc);
306}
307
308static void
309removeExistingServer(bsdpc_t * bsdpc, CFStringRef ServerAddress)
310{
311    int count;
312    int i;
313
314    count = CFArrayGetCount(bsdpc->images);
315    for (i = 0; i < count; i++) {
316	CFDictionaryRef	this_dict = CFArrayGetValueAtIndex(bsdpc->images, i);
317
318	if (CFEqual(ServerAddress,
319		    CFDictionaryGetValue(this_dict, kServerAddress))) {
320	    CFArrayRemoveValueAtIndex(bsdpc->images, i);
321	    return;
322	}
323    }
324    return;
325}
326
327static void
328accumulateImages(bsdpc_t * bsdpc, CFStringRef ServerAddress,
329		 CFNumberRef ServerPriority, CFArrayRef images)
330{
331    CFMutableDictionaryRef	this_dict;
332
333    if (bsdpc->state != bsdpc_state_list) {
334	return;
335    }
336    this_dict = CFDictionaryCreateMutable(NULL, 0,
337					  &kCFTypeDictionaryKeyCallBacks,
338					  &kCFTypeDictionaryValueCallBacks);
339    CFDictionarySetValue(this_dict, kServerAddress, ServerAddress);
340    CFDictionarySetValue(this_dict, kServerPriority, ServerPriority);
341    CFDictionarySetValue(this_dict, kImageList, images);
342    removeExistingServer(bsdpc, ServerAddress);
343    CFArrayAppendValue(bsdpc->images, this_dict);
344    my_CFRelease(&this_dict);
345    if (bsdpc->gathering == FALSE) {
346	struct timeval	t;
347
348	bsdpc->gathering = TRUE;
349	t.tv_sec = 4;
350	t.tv_usec = 0;
351	setTimer(bsdpc, t, gatherDone);
352    }
353    return;
354}
355
356static void
357list_callback(BSDPClientRef client, BSDPClientStatus status,
358	      CFStringRef ServerAddress,
359	      CFNumberRef ServerPriority,
360	      const CFArrayRef list, void *info)
361{
362    bsdpc_t *	bsdpc = (bsdpc_t *)info;
363
364    switch (status) {
365    case kBSDPClientStatusOK:
366	accumulateImages(bsdpc, ServerAddress,
367			 ServerPriority, list);
368	break;
369    case kBSDPClientStatusOperationTimedOut:
370	fprintf(stderr, "No netboot servers found, exiting\n");
371	exit(1);
372	break;
373    default:
374	fprintf(stderr, "List failed, %s\n", BSDPClientStatusString(status));
375	break;
376    }
377    return;
378}
379
380static void
381select_callback(BSDPClientRef client, BSDPClientStatus status, void * info)
382{
383    switch (status) {
384    case kBSDPClientStatusOK:
385	printf("Server confirmed selection\n");
386	exit(0);
387	break;
388    case kBSDPClientStatusServerSentFailure:
389	printf("Server sent failure, selection not confirmed!\n");
390	exit(1);
391	break;
392    default:
393	fprintf(stderr, "Select failed, %s\n", BSDPClientStatusString(status));
394	break;
395    }
396    return;
397}
398
399static void
400select_image(bsdpc_t * bsdpc, int val)
401{
402    CFNumberRef		best_priority = NULL;
403    CFDictionaryRef	best_server_dict = NULL;
404    int 		i;
405    CFDictionaryRef	menu_item;
406    CFArrayRef		served_by;
407    int			served_by_count;
408    BSDPClientStatus	status;
409
410    menu_item = CFArrayGetValueAtIndex(bsdpc->menu, val - 1);
411    served_by = CFDictionaryGetValue(menu_item, kServedBy);
412    served_by_count = CFArrayGetCount(served_by);
413    /* find the "best" server for this image */
414    for (i = 0; i < served_by_count; i++) {
415	CFDictionaryRef	server_dict = CFArrayGetValueAtIndex(served_by, i);
416	CFNumberRef	priority;
417
418	priority = CFDictionaryGetValue(server_dict, kServerPriority);
419	if (best_server_dict == NULL
420	    || (CFNumberCompare(priority, best_priority, NULL)
421		== kCFCompareGreaterThan)) {
422	    best_server_dict = server_dict;
423	    best_priority = priority;
424	}
425    }
426    status
427	= BSDPClientSelect(bsdpc->client,
428			   CFDictionaryGetValue(best_server_dict,
429						kServerAddress),
430			   CFDictionaryGetValue(menu_item,
431						kBSDPImageDescriptionIdentifier),
432			   select_callback, bsdpc);
433    if (status != kBSDPClientStatusOK) {
434	fprintf(stderr, "BSDPClientSelect() failed, %s\n",
435		BSDPClientStatusString(status));
436	exit(1);
437    }
438    bsdpc->state = bsdpc_state_select;
439    return;
440}
441
442static void
443user_input(CFSocketRef s, CFSocketCallBackType type,
444	   CFDataRef address, const void *data, void *info)
445{
446    bsdpc_t * 	bsdpc = (bsdpc_t *)info;
447    char 	choice[128];
448    char 	first_char;
449    char * 	result;
450
451    result = fgets(choice, sizeof(choice), stdin);
452    if (result == NULL) {
453	printf("EOF\n");
454	exit(1);
455    }
456    if (bsdpc->state != bsdpc_state_user_input
457	|| result != choice) {
458	return;
459    }
460    first_char = choice[0];
461    if (first_char >= '1' && first_char <= '9') { /* image selection */
462	int		val = strtoul(choice, 0, 0);
463	if (val >= 1 && val <= CFArrayGetCount(bsdpc->menu)) {
464	    select_image(bsdpc, val);
465	}
466	else {
467	    printf("Value out of range\n");
468	    redisplay(bsdpc);
469	}
470    }
471    else {
472	switch (first_char) {
473	case 'q':
474	    exit(0);
475	    break;
476	default:
477	    printf("Invalid entry\n");
478	    redisplay(bsdpc);
479	    break;
480	}
481    }
482    return;
483}
484
485
486static void
487initialize(const char * ifname, u_int16_t * attrs, int n_attrs)
488{
489    BSDPClientRef	client;
490    CFSocketContext	context = { 0, NULL, NULL, NULL, NULL };
491    CFRunLoopSourceRef	rls = NULL;
492    CFSocketRef		socket = NULL;
493    BSDPClientStatus	status;
494
495    context.info = &bsdpc;
496    socket = CFSocketCreateWithNative(NULL, fileno(stdin), kCFSocketReadCallBack,
497				      user_input, &context);
498    if (socket == NULL) {
499	fprintf(stderr, "CFSocketCreateWithNative failed\n");
500	exit(1);
501    }
502    rls = CFSocketCreateRunLoopSource(NULL, socket, 0);
503    if (rls == NULL) {
504	fprintf(stderr, "CFSocketCreateRunLoopSource failed\n");
505	exit(1);
506    }
507    CFRunLoopAddSource(CFRunLoopGetCurrent(), rls, kCFRunLoopDefaultMode);
508    my_CFRelease(&rls);
509    my_CFRelease(&socket);
510
511    client = BSDPClientCreateWithInterfaceAndAttributes(&status, ifname,
512							attrs, n_attrs);
513    if (client == NULL) {
514	fprintf(stderr, "BSDPClientCreateWithInterface(%s) failed, %s\n",
515		ifname, BSDPClientStatusString(status));
516	BSDPClientFree(&client);
517	exit(1);
518    }
519    status = BSDPClientList(client, list_callback, &bsdpc);
520    if (status != kBSDPClientStatusOK) {
521	fprintf(stderr, "BSDPClientList() failed, %s\n",
522		BSDPClientStatusString(status));
523	BSDPClientFree(&client);
524	exit(1);
525    }
526    bsdpc.state = bsdpc_state_list;
527    bsdpc.client = client;
528    bsdpc.images = CFArrayCreateMutable(NULL, 0, &kCFTypeArrayCallBacks);
529    return;
530}
531
532void
533usage(const char * progname)
534{
535    fprintf(stderr, "usage: %s <options>\n"
536	    "<options> are:\n"
537	    "-i interface : interface name to do BSDP over, default is en0\n"
538	    "-F attrs     : attributes filter\n",
539	    progname);
540    exit(1);
541}
542
543#define MAX_ATTRS	10
544
545int
546main(int argc, char * argv[])
547{
548    int			ch;
549    const char * 	ifname = NULL;
550    u_int16_t		attrs[MAX_ATTRS];
551    int			n_attrs = 0;
552    long		val;
553
554    while ((ch = getopt(argc, argv, "F:i:")) != EOF) {
555	switch (ch) {
556	case 'i':
557	    if (ifname != NULL) {
558		fprintf(stderr, "can't specify interface more than once\n");
559		exit(1);
560	    }
561	    ifname = optarg;
562	    break;
563	case 'F':
564	    errno = 0;
565	    if (n_attrs == MAX_ATTRS) {
566		fprintf(stderr, "too many attributes passed\n");
567		exit(1);
568	    }
569	    val = strtol(optarg, NULL, 0);
570	    if (errno != 0 || val < 0 || val > 65535) {
571		fprintf(stderr, "bad attribute value passed\n");
572		exit(1);
573	    }
574	    attrs[n_attrs++] = val;
575	    break;
576	default:
577	    break;
578	}
579    }
580    if ((argc - optind) != 0) {
581	usage(argv[0]);
582    }
583    if (ifname == NULL) {
584	ifname = "en0";
585    }
586
587    printf("Discovering NetBoot servers...\n");
588    if (n_attrs == 0) {
589	initialize(ifname, NULL, 0);
590    }
591    else {
592	initialize(ifname, attrs, n_attrs);
593    }
594    CFRunLoopRun();
595    printf("CFRunLoop done\n");
596    exit(0);
597}
598