fix perms
[mpi-image.git] / scripts / init.py
1 #!/usr/bin/python
2
3 import sys
4 import os
5 import signal
6 import socket
7 import fcntl
8 import pdb
9 import getpass
10
11 INSANE = 500
12 BLOCKING_DELAY = 30
13
14 HTML_SOURCE = "www.vicci.org/files/mpi/"
15 MPI_COPY_URL = HTML_SOURCE + "mpi-copy.py"
16 GENERATE_HOSTFILE_URL = HTML_SOURCE + "generate-hostfile.py"
17 SSHD_CONFIG_URL = HTML_SOURCE + "sshd_config"
18 ORTED_URL = HTML_SOURCE + "orted"
19 BIND_PUBLIC_URL = HTML_SOURCE + "bind_public.so"
20 MPIRUN_URL = HTML_SOURCE + "mpirun"
21
22 def alarm(signum, junk):
23     # Do nothing, default implementation does needful
24     return
25
26 def fetch_output(infd,outfile):
27     try:
28         os.write(infd,'1')
29         output = ''
30
31         while (True):
32             batch = outfile.read()
33             if (batch):
34                 output+=batch
35             else:
36                 break
37     except IOError,e:
38         print "IO Error Accessing Vsys"
39         raise e
40
41     return output
42
43 def fetch_output_nofail(inpath,outpath):
44     pid = os.fork()
45     
46     output = None
47
48     if (pid==0):
49         outfd = os.open(outpath, os.O_RDONLY|os.O_NONBLOCK)
50         infd = os.open(inpath, os.O_WRONLY)
51         #infile = os.fdopen(infd,'w')
52
53         flags = fcntl.fcntl(outfd, fcntl.F_GETFL)
54         fcntl.fcntl(outfd, fcntl.F_SETFL, flags & ~os.O_NONBLOCK)
55         outfile = os.fdopen(outfd)
56
57         output = fetch_output(infd, outfile)
58     else:
59         signal.alarm(BLOCKING_DELAY)
60         try:
61             os.waitpid(pid,0)
62             exit(1)
63         except OSError:
64             ## Failed
65             pass
66
67     return output
68
69
70 ### main()
71
72 os.system("sudo mkdir -m 777 /usr/local/mpifix")
73
74 os.system("wget -q -O /usr/bin/mpi-copy.py " + MPI_COPY_URL)
75 os.system("wget -q -O /usr/bin/generate-hostfile.py " + GENERATE_HOSTFILE_URL)
76 os.system("wget -q -O /tmp/sshd_config " + SSHD_CONFIG_URL)
77 os.system("wget -q -O /usr/local/mpifix/orted " + ORTED_URL)
78 os.system("wget -q -O /usr/local/mpifix/bind_public.so " + BIND_PUBLIC_URL)
79 os.system("wget -q -O /usr/local/mpifix/mpirun " + MPIRUN_URL)
80
81 os.system("chmod 755 /usr/bin/mpi-copy.py")
82 os.system("chmod 755 /usr/bin/generate-hostfile.py")
83 os.system("chmod 755 /usr/local/mpifix/orted")
84 os.system("chmod 755 /usr/local/mpifix/mpirun")
85
86 for fn in ["/usr/bin/mpi-copy.py", "/usr/bin/generate-hostfile.py", "/tmp/sshd_config"]:
87     if not os.path.exists(fn):
88         print "missing", fn
89         sys.exit(-1)
90     if os.path.getsize(fn)==0:
91         print "zero byte file", fn
92         sys.exit(-1)
93
94 f = open("/etc/slicename", "r")
95 homedir = "/home/" + f.read()
96 homedir = homedir.strip()
97 s = "mkdir " + homedir + "/.ssh"
98 os.system(s)
99
100 outpath = '/vsys/%s.out'%sys.argv[1]
101 inpath = '/vsys/%s.in'%sys.argv[1]
102
103 ###
104 ### If vsys is not set up, wait a bit and try again
105
106 try:
107     counter = 1
108     while True:
109             signal.signal(signal.SIGALRM, alarm)
110             output = fetch_output_nofail(inpath,outpath)
111             if (output is not None or counter==INSANE):
112                 break
113             print "No access to vsys. Looping (%d)..."%counter
114             counter = counter+1
115
116     if (output is None):
117         print "Could not connect to Vsys. Giving up."
118         exit(1)
119
120     # get the ip address of the host, we'll need it when setting up sshd
121     hostname = socket.gethostname()
122     hostipaddr = socket.gethostbyname(hostname)
123     print "found hostname", hostname, "and ip", hostipaddr
124
125     wholestring = output.split("\n")
126     seenSSHKey = False
127     for eachline in wholestring:
128       if eachline == "\n" or eachline == "":
129         continue
130       line = eachline.split()
131       if line[0] == "vsys_sshKey:":
132         pathname = homedir + "/.ssh/id_rsa"
133         g = open(pathname, "w+")
134         s = line[1] + " " + line[2] + " " + line[3] + " " + line[4] + "\n"
135     #    for i in line[5:len(line)-4]:
136     #      str = str + i + "\n"
137     #    for i in line[len(line)-4:]:
138     #      str = str + i + " "
139     #    str = str + "\n"
140         g.write(s)
141         seenSSHKey = True
142       elif seenSSHKey:
143         if line[len(line)-1] == "KEY-----":
144           s = line[0] + " " + line[1] + " " + line[2] + " " + line[3] + "\n"
145           g.write(s)
146           seenSSHKey = False
147           g.close()
148         else:
149           g.write(line[0] + "\n")
150       elif line[0] == "vsys_portNumber:":
151         pathname = homedir + "/.ssh/config"
152         g = open(pathname, "w+")
153         s = "Port " + line[1] + "\n"
154         g.write(s)
155         g.write("StrictHostKeyChecking no\n")
156         g.write("UserKnownHostsFile /dev/null\n")
157         g.write("LogLevel quiet\n")
158         g.close()
159         h = open("/tmp/portFile", "w+")
160         h.write(s)
161         h.close()
162         os.system("sudo rm /etc/ssh/sshd_config")
163         os.system("sudo cat /tmp/sshd_config /tmp/portFile > /tmp/newsshd_config")
164         os.system('sudo echo "ListenAddress ' + hostipaddr + '" >> /tmp/newsshd_config')
165         os.system("sudo mv /tmp/newsshd_config /etc/ssh/sshd_config")
166       elif line[0] == "vsys_sshKey.pub:":
167         pathname = homedir + "/.ssh/id_rsa.pub"
168         g = open(pathname, "w+")
169         s = ""
170         for i in line[1:]:
171           s = s  + i + " "
172         s = s.strip(" ")
173         s = s + "\n"
174         g.write(s)
175         g.close()
176         dest = homedir + "/.ssh/authorized_keys"
177         os.system("cp " + pathname + " " + dest)
178
179     username = homedir[6:]
180     s = "sudo chown -R " + username + " " + homedir + "/.ssh"
181     os.system(s)
182     s = "sudo chgrp -R slices " + homedir + "/.ssh"
183     os.system(s)
184 #    s = "sudo chown " + username + " " + homedir + "/mpi-copy.py"
185 #    os.system(s)
186 #    s = "sudo chgrp slices " + homedir + "/mpi-copy.py"
187 #    os.system(s)
188 #    s = "sudo chmod u+x " + homedir + "/mpi-copy.py"
189 #    os.system(s)
190
191     # /etc/profile isn't working for noninteractive ssh sessions, so set the mpi
192     # paths in .bashrc
193     open(homedir+"/.bashrc","w").write("export PATH=/usr/local/mpifix:$PATH:/usr/lib64/openmpi/bin\nexport LD_LIBRARY_PATH=/usr/lib64/openmpi/lib:$LD_LIBRARY_PATH\n")
194
195     os.system("chmod og-rw " + homedir + "/.ssh/id_rsa")
196     os.system("chmod og-rw " + homedir + "/.ssh/id_rsa.pub")
197     os.system("chmod og-rw " + homedir + "/.ssh/config")
198     os.system("rm -rf /tmp/portFile")
199     os.system("sudo rm -rf /tmp/sshd_config")
200 except Exception,e:
201     print "Unexpected error:", sys.exc_info()[0]
202
203