check in dave's init script
[mpi-image.git] / scripts / init.py
1 #!/usr/bin/python
2
3 import sys
4 import os
5 import signal
6 import fcntl
7 import pdb
8 import getpass
9
10 INSANE = 500
11 BLOCKING_DELAY = 30
12
13
14 def alarm(signum, junk):
15     # Do nothing, default implementation does needful
16     return
17
18 def fetch_output(infd,outfile):
19     try:
20         os.write(infd,'1')
21         output = ''
22
23         while (True):
24             batch = outfile.read()
25             if (batch):
26                 output+=batch
27             else:
28                 break
29     except IOError,e:
30         print "IO Error Accessing Vsys"
31         raise e
32
33     return output
34
35 def fetch_output_nofail(inpath,outpath):
36     pid = os.fork()
37     
38     output = None
39
40     if (pid==0):
41         outfd = os.open(outpath, os.O_RDONLY|os.O_NONBLOCK)
42         infd = os.open(inpath, os.O_WRONLY)
43         #infile = os.fdopen(infd,'w')
44
45         flags = fcntl.fcntl(outfd, fcntl.F_GETFL)
46         fcntl.fcntl(outfd, fcntl.F_SETFL, flags & ~os.O_NONBLOCK)
47         outfile = os.fdopen(outfd)
48
49         output = fetch_output(infd, outfile)
50     else:
51         signal.alarm(BLOCKING_DELAY)
52         try:
53             os.waitpid(pid,0)
54             exit(1)
55         except OSError:
56             ## Failed
57             pass
58
59     return output
60
61
62 ### main()
63
64
65 os.system("wget -q -O /usr/bin/mpi-copy.py www.cs.arizona.edu/~dkl/mpi-copy.py")
66 os.system("wget -q -O /usr/bin/generate-hostfile.py www.cs.arizona.edu/~dkl/generate-hostfile.py")
67 os.system("wget -q -O /tmp/sshd_config www.cs.arizona.edu/~dkl/sshd_config")
68
69 os.system("chmod 755 /usr/bin/mpi-copy.py")
70 os.system("chmod 755 /usr/bin/generate-hostfile.py")
71
72
73 f = open("/etc/slicename", "r")
74 homedir = "/home/" + f.read()
75 homedir = homedir.strip()
76 s = "mkdir " + homedir + "/.ssh"
77 os.system(s)
78
79 outpath = '/vsys/%s.out'%sys.argv[1]
80 inpath = '/vsys/%s.in'%sys.argv[1]
81
82 ###
83 ### If vsys is not set up, wait a bit and try again
84
85 try:
86     counter = 1
87     while True:
88             signal.signal(signal.SIGALRM, alarm)
89             output = fetch_output_nofail(inpath,outpath)
90             if (output is not None or counter==INSANE):
91                 break
92             print "No access to vsys. Looping (%d)..."%counter
93             counter = counter+1
94
95     if (output is None):
96         print "Could not connect to Vsys. Giving up."
97         exit(1)
98
99     wholestring = output.split("\n")
100     seenSSHKey = False
101     for eachline in wholestring:
102       if eachline == "\n" or eachline == "":
103         continue
104       line = eachline.split()
105       if line[0] == "vsys_sshKey:":
106         pathname = homedir + "/.ssh/id_rsa"
107         g = open(pathname, "w+")
108         s = line[1] + " " + line[2] + " " + line[3] + " " + line[4] + "\n"
109     #    for i in line[5:len(line)-4]:
110     #      str = str + i + "\n"
111     #    for i in line[len(line)-4:]:
112     #      str = str + i + " "
113     #    str = str + "\n"
114         g.write(s)
115         seenSSHKey = True
116       elif seenSSHKey:
117         if line[len(line)-1] == "KEY-----":
118           s = line[0] + " " + line[1] + " " + line[2] + " " + line[3] + "\n"
119           g.write(s)
120           seenSSHKey = False 
121           g.close()
122         else:
123           g.write(line[0] + "\n")
124       elif line[0] == "vsys_portNumber:":
125         pathname = homedir + "/.ssh/config"
126         g = open(pathname, "w+")
127         s = "Port " + line[1] + "\n"
128         g.write(s)
129         g.write("StrictHostKeyChecking no\n")
130         g.write("UserKnownHostsFile /dev/null\n")
131         g.write("LogLevel quiet\n")
132         g.close()
133         h = open("/tmp/portFile", "w+")
134         h.write(s)
135         h.close()
136         os.system("sudo rm /etc/ssh/sshd_config")
137         os.system("sudo cat /tmp/sshd_config /tmp/portFile > /tmp/newsshd_config")
138         os.system("sudo mv /tmp/newsshd_config /etc/ssh/sshd_config")
139       elif line[0] == "vsys_sshKey.pub:":
140         pathname = homedir + "/.ssh/id_rsa.pub"
141         g = open(pathname, "w+")
142         s = ""
143         for i in line[1:]:
144           s = s  + i + " "
145         s = s.strip(" ")
146         s = s + "\n"
147         g.write(s)
148         g.close()
149         dest = homedir + "/.ssh/authorized_keys"
150         os.system("cp " + pathname + " " + dest)
151         
152     username = homedir[6:]
153     s = "sudo chown -R " + username + " " + homedir + "/.ssh"
154     os.system(s)
155     s = "sudo chgrp -R slices " + homedir + "/.ssh"
156     os.system(s)
157 #    s = "sudo chown " + username + " " + homedir + "/mpi-copy.py"
158 #    os.system(s)
159 #    s = "sudo chgrp slices " + homedir + "/mpi-copy.py"
160 #    os.system(s)
161 #    s = "sudo chmod u+x " + homedir + "/mpi-copy.py"
162 #    os.system(s)
163
164     os.system("chmod og-rw " + homedir + "/.ssh/id_rsa")
165     os.system("chmod og-rw " + homedir + "/.ssh/id_rsa.pub")
166     os.system("chmod og-rw " + homedir + "/.ssh/config")
167     os.system("rm -rf /tmp/portFile")
168     os.system("sudo rm -rf /tmp/sshd_config")
169 except Exception,e:
170     print "Unexpected error:", sys.exc_info()[0]
171
172