Small restructure

Signed-off-by: JasonMillward <jason@jcode.me>
Added try catch to cleanup function