diff --git a/firestarter b/firestarter
index ad4c7d5..bc964f9 100755
--- a/firestarter
+++ b/firestarter
@@ -55,23 +55,6 @@ startfile() {
 	fi
 }
 
-# Traps
-trap_exit() {
-	log "Logging out"
-	if command -v loginctl > /dev/null 2>&1; then
-		# Use loginctl if possible
-		local sessionid="$(< /proc/self/sessionid)"
-		if [ -n "$sessionid" ]; then
-			loginctl terminate-session "$sessionid"
-		fi
-	else
-		# Otherwise just brute it
-		kill $(jobs -p)
-	fi
-	return 0
-}
-trap trap_exit EXIT
-
 # Steps in execution
 step_generate() {
 	log "Creating default config setup in \"$g_configdir\""
@@ -399,9 +382,24 @@ step_wait() {
 	[ "$dryrun" = "1" ] && exit 0
 	log "Waiting for programs to exit"
 	log "Send any termination signal to firestarter to log out"
+	trap step_logout EXIT
 	wait
 	exit 0
 }
+step_logout() {
+	log "Logging out"
+	if command -v loginctl > /dev/null 2>&1; then
+		# Use loginctl if possible
+		local sessionid="$(< /proc/self/sessionid)"
+		if [ -n "$sessionid" ]; then
+			loginctl terminate-session "$sessionid"
+		fi
+	else
+		# Otherwise just brute it
+		kill $(jobs -p)
+	fi
+	return 0
+}
 
 # Main
 main() {