diff --git a/target/linux/au1000/base-files/etc/diag.sh b/target/linux/au1000/base-files/etc/diag.sh
index d97461532a..65e3011be5 100644
--- a/target/linux/au1000/base-files/etc/diag.sh
+++ b/target/linux/au1000/base-files/etc/diag.sh
@@ -1,25 +1,21 @@
 #!/bin/sh
-# Copyright (C) 2007 OpenWrt.org
+# Copyright (C) 2007-2013 OpenWrt.org
 
-set_led() {
-	local led="$1"
-	local state="$2"
-	[ -d "/sys/class/leds/mtx1:$led" ] && echo "$state" > "/sys/class/leds/mtx1:$led/brightness"
-}
+. /lib/functions/leds.sh
 
 set_state() {
 	case "$1" in
-		preinit)
-			set_led green 0
-			set_led red 1
+	preinit)
+		led_off "mtx1:green"
+		led_on "mtx1:red"
 		;;
-		failsafe)
-			set_led green 1
-			set_led red 1
+	failsafe)
+		led_on "mtx1:green"
+		led_on "mtx1:red"
 		;;
-		done)
-			set_led green 1
-			set_led red 0
+	done)
+		led_on "mtx1:green"
+		led_off "mtx1:red"
 		;;
 	esac
 }