Updated python script to work in linux