Fixed code to a working point.
Signed-off-by: James Seiters <jamesseiters@gmail.com>