Improved CLI
Added lacking import