Made script a bit more robust